Skip to content
Snippets Groups Projects
Commit a850ba43 authored by Tim Tobias Bauerle's avatar Tim Tobias Bauerle
Browse files

Fixed cuda error by loading AE weights explicitly on CPU

parent 10bb5fd2
No related branches found
No related tags found
1 merge request!21Fixed AE weights loading issue
......@@ -5,7 +5,7 @@ from torch import nn
class SerializedModel(nn.Module):
def __init__(self, model: nn.Module, path: str):
super().__init__()
model.load_state_dict(torch.load(path))
model.load_state_dict(torch.load(path, map_location="cpu"))
self.model = model
def forward(self, x):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment