I have a weird problem. I can load the state_dict in my machine which was trained with Pytorch lightning on the GPU. However if i push it to docker it does not load. The model is a Pytorch Lightning module.

On Machine
loaded_model = LitModel()
loaded_model.load_state_dict(torch.load(‘models/trained-model-2021.pt’), strict=False)
Parameter containing:
tensor([[ 0.0817, -0.2579, -0.0080,  …,  0.1619, -0.0783, -0.0087],
[ 0.0866, -0.2708,  0.2528,  …,  0.0203, -0.2319,  0.1011],
[ 0.0564,  0.0195, -0.2858,  …, -0.1899,  0.1004, -0.1013],
…,
[ 0.2584, -0.2400, -0.0358,  …, -0.1391,  0.2278,  0.0264],
[ 0.1901, -0.3056, -0.2537,  …, -0.0747,  0.2703,  0.0803],
[-0.2639, -0.0381,  0.0538,  …,  0.2504,  0.2152,  0.1084]],
requires_grad=True)

On Docker
device = torch.device(‘cpu’)
loaded_model = LitModel()
loaded_model.load_state_dict(torch.load(‘models/trained-model-2021.pt’, map_location=device), strict=False)
loaded_model.load_state_dict(torch.load(‘src/trained-model-2021-11-23–23:08-hs.pt’, map_location=device), strict=True)
Traceback (most recent call last):
File “”, line 1, in 
File “/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py”, line 1407, in load_state_dict
self.class.name, “\n\t”.join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for LitModel:
Missing key(s) in state_dict: “layer_1.weight”, “layer_1.bias”, “layer_3.weight”, “layer_3.bias”.
Unexpected key(s) in state_dict: “model”.
Any ideas what could be wrong here.