Fixes to prep for weights_only default flip (#2514)

Summary:
Some fixes for pytorch/pytorch#137602


Reviewed By: xuzhao9

Differential Revision: D64628614

Pulled By: mikaylagawarecki

torchbenchmark/models/functorch_maml_omniglot/__init__.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 from typing import Tuple
 
+import numpy as np
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -73,7 +75,16 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
         self.model = net
 
         root = str(Path(__file__).parent.parent)
-        self.meta_inputs = torch.load(f"{root}/maml_omniglot/batch.pt")
+        with torch.serialization.safe_globals(
+            [
+                np.core.multiarray._reconstruct,
+                np.ndarray,
+                np.dtype,
+                np.dtypes.Float32DType,
+                np.dtypes.Int64DType,
+            ]
+        ):
+            self.meta_inputs = torch.load(f"{root}/maml_omniglot/batch.pt")
         self.meta_inputs = tuple(
             [torch.from_numpy(i).to(self.device) for i in self.meta_inputs]
         )

torchbenchmark/models/maml_omniglot/__init__.py

@@ -20,6 +20,7 @@
 from typing import Tuple
 
 import higher
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -79,7 +80,16 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
         self.model = net
 
         root = str(Path(__file__).parent)
-        self.meta_inputs = torch.load(f"{root}/batch.pt")
+        with torch.serialization.safe_globals(
+            [
+                np.core.multiarray._reconstruct,
+                np.ndarray,
+                np.dtype,
+                np.dtypes.Float32DType,
+                np.dtypes.Int64DType,
+            ]
+        ):
+            self.meta_inputs = torch.load(f"{root}/batch.pt")
         self.meta_inputs = tuple(
             [torch.from_numpy(i).to(self.device) for i in self.meta_inputs]
         )

torchbenchmark/models/opacus_cifar10/__init__.py

@@ -1,3 +1,4 @@
+import os
 from typing import Tuple
 
 import torch
@@ -29,7 +30,13 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
         )
 
         self.model = models.resnet18(num_classes=10)
+        prev_wo_envvar = os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", None)
+        os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
         self.model = ModuleValidator.fix(self.model)
+        if prev_wo_envvar is None:
+            del os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"]
+        else:
+            os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = prev_wo_envvar
         self.model = self.model.to(device)
 
         # Cifar10 images are 32x32 and have 10 classes