Spaces:

cherubicxn
/

ScaleLSD

Running on Zero

cherubicxn commited on Jun 13

Commit

e33fa08

1 Parent(s): 1828c29

update

Files changed (1) hide show

scalelsd/ssl/misc/train_utils.py CHANGED Viewed

@@ -46,7 +46,7 @@ def load_scalelsd_model(ckpt_path, device='cuda'):
     model = ScaleLSD(gray_scale=True, use_layer_scale=use_layer_scale)
     model = model.eval().to(device)
-    state_dict = torch.load(ckpt_path, map_location='cpu',weights_only=False)
     try:
         model.load_state_dict(state_dict['model_state'])
     except:

     model = ScaleLSD(gray_scale=True, use_layer_scale=use_layer_scale)
     model = model.eval().to(device)
+    state_dict = torch.load(ckpt_path, map_location='cpu', weights_only=True)
     try:
         model.load_state_dict(state_dict['model_state'])
     except: