Spaces:

patrickjdarrow
/

oam_autoencoder

Sleeping

patrickjdarrow commited on Oct 8, 2023

Commit

70d8d45

1 Parent(s): b8f5ef0

Fixed model load scheme

Files changed (1) hide show

app.py CHANGED Viewed

@@ -8,7 +8,7 @@ from autoencoder import Autoencoder
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 model = Autoencoder()
-model.load_state_dict('model.pt', map_location=device)
 resize = Resize((224))
 to_tensor = ToTensor()

 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 model = Autoencoder()
+model.load_state_dict(torch.load('model.pt', map_location=device))
 resize = Resize((224))
 to_tensor = ToTensor()