Spaces:

LuxExistentia
/

Age-PT

Runtime error

LuxExistentia commited on Oct 6, 2024

Commit

f1eecad

1 Parent(s): 983448b

Fixed a bug with loading weights to cpu device

Files changed (1) hide show

app.py CHANGED Viewed

@@ -26,7 +26,7 @@ model_classifier = nn.Sequential(nn.Linear(512, 512),
                                 nn.GELU(),
                                 nn.Linear(512, 1))
 model = nn.Sequential(model, model_classifier)
-model.load_state_dict(state_dict=torch.load(FILE_NAME, weights_only=True))
 test_transform = timm.data.create_transform(input_size=224, is_training=False, interpolation="bicubic")

                                 nn.GELU(),
                                 nn.Linear(512, 1))
 model = nn.Sequential(model, model_classifier)
+model.load_state_dict(state_dict=torch.load(FILE_NAME, weights_only=True, map_location=torch.device(DEVICE)))
 test_transform = timm.data.create_transform(input_size=224, is_training=False, interpolation="bicubic")