The PyTorch model won't load in PyTorch 0.4.*

junyanz / interactive-deep-colorization

Deep learning software for colorizing black and white images with a few clicks.

https://richzhang.github.io/ideepcolor/

MIT License

2.69k stars 447 forks source link

The PyTorch model won't load in PyTorch 0.4.* #54

Closed l3robot closed 5 years ago

l3robot commented 6 years ago

Hi, It is quite a small details, though useful to know. The pre-trained model won't work in PyTorch 0.4.* because normalization layers now track the number of batches with the tensor num_batches_tracked [see code]. When loading a state_dict, it then expects these tensors and crashes on a KeyError. Maybe just adding the pytorch version in requirements could help future users. Thanks

junyanz commented 6 years ago

Cool. I updated the README and add PyTorch 0.5. (I added a patch to the network loading code just in case people are still using PyTorch 0.4. )

l3robot commented 5 years ago

Excellent! thanks sry for the delay