Github Actions setup

Author: RohitXJ

.github/workflows/test-export.yml

@@ -1,6 +1,8 @@
 name: Test Few-Shot Model Export
 
 on:
+  workflow_dispatch:
+    
   push:
     branches: [ main ]
   pull_request:

testing/test1.py

@@ -1,14 +1,29 @@
+# testing/test1.py
+
 import torch
+import torchvision.transforms as T
+import pprint
+
+# Trust the transforms used in the model export
+torch.serialization.add_safe_globals({
+    'torchvision.transforms.transforms.Compose': T.Compose,
+    'torchvision.transforms.transforms.Resize': T.Resize,
+    'torchvision.transforms.transforms.ToTensor': T.ToTensor
+})
 
 def test_model_export(model_path):
-    print(f"Loading model from: {model_path}")
-    config = torch.load(model_path, map_location='cpu')
+    print(f"\nLoading model from: {model_path}")
+    
+    config = torch.load(model_path, map_location='cpu', weights_only=False)
+
+    print("\n--- Exported Model Content ---")
 
     for key, value in config.items():
-        if torch.is_tensor(value):
-            print(f"{key}: Tensor with shape {tuple(value.shape)}")
+        if isinstance(value, torch.Tensor):
+            print(f"{key}: Tensor shape = {tuple(value.shape)}")
         else:
-            print(f"{key}: {value}")
+            print(f"{key}:")
+            pprint.pprint(value)
 
-# Run
-test_model_export("export/fewshot_model.pt")
+if __name__ == "__main__":
+    test_model_export("export/fewshot_model.pt")