update

scalelsd/ssl/misc/train_utils.py

@@ -42,7 +42,13 @@ def fix_seeds(random_seed):
 
 def load_scalelsd_model(ckpt_path, device='cuda'):
     """load model"""
-    use_layer_scale = False if 'v1' in ckpt_path else True
+    use_layer_scale = True
+    if os.path.basename(ckpt_path) == 'scalelsd-vitbase-v1-train-sa1b.pt':
+        use_layer_scale = False
+    elif os.path.basename(ckpt_path) == 'scalelsd-vitbase-v2-train-sa1b.pt':
+        use_layer_scale = True
+    else:
+        raise ValueError(f'Unknown model: {os.path.basename(ckpt_path)}')
 
     model = ScaleLSD(gray_scale=True, use_layer_scale=use_layer_scale)
     model = model.eval().to(device)