-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
🐛 Bug
I noticed that uninitialised module buffers don't show up in what state_dict() returns. Is this a bug or a feature?
This makes loading from checkpoints a bit difficult.
To Reproduce
import torch.nn as nn
import torch
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("some_state", None)
module_original = MyModule()
module_original.some_state = torch.tensor(42)
# later when loading checkpoint
module_copy = MyModule()
module_copy.load_state_dict(module_original.state_dict())Fails with:
Traceback (most recent call last):
File "bug.py", line 17, in <module>
module_copy.load_state_dict(module_original.state_dict())
File "/home/milad/miniconda3/envs/pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1045, in load_state_dict
self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for MyModule:
Unexpected key(s) in state_dict: "some_state".
Environment
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A
OS: Manjaro Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: Could not collect
CMake version: version 3.18.3
Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce GTX 1660 SUPER
Nvidia driver version: 440.100
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] pytorch-ignite==0.4.0.post1
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] mkl 2020.2 256
[conda] mkl-service 2.3.0 py37he904b0f_0
[conda] mkl_fft 1.2.0 py37h23d657b_0
[conda] mkl_random 1.1.1 py37h0573a6f_0
[conda] numpy 1.19.1 py37hbc911f0_0
[conda] numpy-base 1.19.1 py37hfa32c7d_0
[conda] pytorch 1.6.0 py3.7_cuda10.2.89_cudnn7.6.5_0 pytorch
[conda] pytorch-ignite 0.4.0.post1 pypi_0 pypi
[conda] torchvision 0.7.0 py37_cu102 pytorch