@@ -14,3 +14,1 @@
-    from torch.onnx import operators
-
-    return operators.shape_as_tensor(image)[-2:]
+    # from torch.onnx import operators

@@ -64,2 +67,4 @@
         for k in to_load:
             if sys.modules.get(k) is not None:
+                if self.enabled:
+                    self.torch_unique[k] = sys.modules.get(k)