Update generate_full_weights_1024x1024.py

Alias-z · Mar 19, 2024 · a90ec5f · a90ec5f
1 parent d4535c6
commit a90ec5f
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tools/generate_full_weights_1024x1024.py b/tools/generate_full_weights_1024x1024.py
@@ -38,10 +38,11 @@ def main(args):
     backbone = args.backbone
     rein_head = args.rein_head
 
-    if not osp.isfile(dinov2_segmentor_path):
-        weight = torch.load(rein_head, map_location='cpu')
-        weight['state_dict'].update({f'backbone.{k}': v for k, v in load_backbone(backbone).items()})
-        torch.save(weight, dinov2_segmentor_path)
+    weight = torch.load(rein_head, map_location='cpu')
+    if 'state_dict' not in weight:
+        weight=dict(state_dict=weight)
+    weight['state_dict'].update({f'backbone.{k}': v for k, v in load_backbone(backbone).items()})
+    torch.save(weight, dinov2_segmentor_path)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Load and process model weights.")