fix: making sure elements in pytorch files are tensors

dreadnode · Nov 8, 2024 · c40b642 · c40b642
1 parent 83961ea
commit c40b642
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/core/handlers/pytorch/inspect.py b/src/core/handlers/pytorch/inspect.py
@@ -2,6 +2,7 @@
 import json
 import os
 import argparse
+import numpy as np
 
 
 def main():
@@ -51,6 +52,13 @@ def main():
         model = model["model"]
 
     for tensor_name, tensor in model.items():
+        # make sure it's a tensor
+        if not isinstance(tensor, torch.Tensor):
+            try:
+                tensor = torch.tensor(tensor)
+            except:
+                continue
+
         inspection["data_size"] += tensor.shape.numel() * tensor.element_size()
 
         shape = list(tensor.shape)