[BUG] Pre-trained model unloadable

[jdeschamps]

Describe the bug In the bbbc010-2012 predict notebook, the default cell loading the pre-trained model fails.

It seems it cannot load the state_dict because the module path has changed (encoder.initial_block.conv.weight -> module.encoder.initial_block.conv.weight):

File [/localscratch/miniconda3/envs/EmbedSeg/lib/python3.10/site-packages/torch/nn/modules/module.py:2152](https://vscode-remote+ssh-002dremote-002bvdi8.vscode-resource.vscode-cdn.net/localscratch/miniconda3/envs/EmbedSeg/lib/python3.10/site-packages/torch/nn/modules/module.py:2152), in Module.load_state_dict(self, state_dict, strict, assign)
   2147         error_msgs.insert(
   2148             0, 'Missing key(s) in state_dict: {}. '.format(
   2149                 ', '.join(f'"{k}"' for k in missing_keys)))
   2151 if len(error_msgs) > 0:
-> 2152     raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   2153                        self.__class__.__name__, "\n\t".join(error_msgs)))
   2154 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for BranchedERFNet:
	Missing key(s) in state_dict: "encoder.initial_block.conv.weight", "encoder.initial_block.conv.bias", [...]
	Unexpected key(s) in state_dict: "module.encoder.initial_block.conv.weight", "module.encoder.initial_block.conv.bias", [...]

(I shorten the error, obviously it lists all the layers)

I used pytorch 2.1.0