Spaces:

mjavadmt
/

EnergySmell

Build error

mjavadmt commited on Mar 27

Commit

da5c744

1 Parent(s): d407fa8

change location

Files changed (1) hide show

model.py CHANGED Viewed

@@ -31,6 +31,6 @@ class EnergySmellsDetector(nn.Module):
         # Load model
         model = EnergySmellsDetector(PRETRAINED_MODEL)
-        model.load_state_dict(torch.load(model_path))
         return model

         # Load model
         model = EnergySmellsDetector(PRETRAINED_MODEL)
+        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
         return model