Spaces:
Sleeping
Sleeping
Pavel Malov
commited on
Commit
·
5eae5da
1
Parent(s):
28f6ce1
CPU only
Browse files- inference.py +1 -1
inference.py
CHANGED
|
@@ -27,7 +27,7 @@ class ClassificationHead(nn.Module):
|
|
| 27 |
|
| 28 |
self.nnet = Nnet()
|
| 29 |
|
| 30 |
-
ckpt = torch.load("resources/model.ckpt")
|
| 31 |
self.nnet.load_state_dict(ckpt['state_dict'], strict=False)
|
| 32 |
|
| 33 |
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
| 27 |
|
| 28 |
self.nnet = Nnet()
|
| 29 |
|
| 30 |
+
ckpt = torch.load("resources/model.ckpt", map_location=torch.device('cpu'))
|
| 31 |
self.nnet.load_state_dict(ckpt['state_dict'], strict=False)
|
| 32 |
|
| 33 |
def forward(self, x: torch.Tensor) -> torch.Tensor:
|