-
Notifications
You must be signed in to change notification settings - Fork 117
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Cannot load the model when trained with batch_size == 1 #140
Comments
Additional information that I just found: Otherwise, do you think adding track_running_stats=False to all BatchNorm2d would work? Maybe replacing BatchNorm2d by InstanceNorm2d? |
Yeah, when |
BTW, I mean loading the each |
Hi @ZhengPeng7
Thank you so much for your amazing work!
To verify that the training goes well, I made a small change to the train_epoch function as:
The exported model seems to be slightly lighter than your trained model (yours is 885.1 MB while this one got 884.8 MB).
When I load the model using this snippet:
It raised an error of:
I assume my code is missing something, can you please guide me? Thank you so much!
The text was updated successfully, but these errors were encountered: