🚀[FEA]: add torchfix linting

[Ethan Wu]

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Medium

Please provide a clear description of problem you would like to solve.

I ran torchfix on this repo and got the following message.

~/earth2mip$ torchfix .
earth2mip/beta/perturbation/bv.py:104:17: TOR101 Use of deprecated function torch.norm
earth2mip/beta/perturbation/bv.py:104:33: TOR101 Use of deprecated function torch.norm
build/lib/earth2mip/diagnostic/climate_net.py:556:31: TOR102 [*] `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
earth2mip/diagnostic/climate_net.py:556:31: TOR102 [*] `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
build/lib/earth2mip/ensemble_utils.py:229:13: TOR101 Use of deprecated function torch.norm
build/lib/earth2mip/ensemble_utils.py:229:29: TOR101 Use of deprecated function torch.norm
earth2mip/ensemble_utils.py:231:13: TOR101 Use of deprecated function torch.norm
earth2mip/ensemble_utils.py:231:29: TOR101 Use of deprecated function torch.norm
--- /home/ethan/earth2mip/earth2mip/diagnostic/climate_net.py
+++ /home/ethan/earth2mip/earth2mip/diagnostic/climate_net.py
@@ -551,11 +551,11 @@
         model = CGNetModule(
             channels=len(IN_CHANNELS),
             classes=len(OUT_CHANNELS),
         )
         weights_path = package.get("weights.tar")
-        model.load_state_dict(torch.load(weights_path, map_location=device))
+        model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
         model.eval()

         input_center = torch.Tensor(np.load(package.get("global_means.npy")))[
             :, None, None
         ]
--- /home/ethan/earth2mip/build/lib/earth2mip/diagnostic/climate_net.py
+++ /home/ethan/earth2mip/build/lib/earth2mip/diagnostic/climate_net.py
@@ -551,11 +551,11 @@
         model = CGNetModule(
             channels=len(IN_CHANNELS),
             classes=len(OUT_CHANNELS),
         )
         weights_path = package.get("weights.tar")
-        model.load_state_dict(torch.load(weights_path, map_location=device))
+        model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
         model.eval()

         input_center = torch.Tensor(np.load(package.get("global_means.npy")))[
             :, None, None
         ]
build/lib/earth2mip/networks/fcnv2_sm.py:155:15: TOR102 [*] `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
--- /home/ethan/earth2mip/build/lib/earth2mip/networks/fcnv2_sm.py
+++ /home/ethan/earth2mip/build/lib/earth2mip/networks/fcnv2_sm.py
@@ -150,11 +150,11 @@

     local_center = np.load(package.get("global_means.npy"))
     local_std = np.load(package.get("global_stds.npy"))

     weights_path = package.get("weights.tar")
-    weights = torch.load(weights_path, map_location=device)
+    weights = torch.load(weights_path, map_location=device, weights_only=True)
     fixed_weights = _fix_state_dict_keys(weights["model_state"], add_module=False)
     core_model.load_state_dict(fixed_weights)

     grid = earth2mip.grid.equiangular_lat_lon_grid(721, 1440)
     dt = datetime.timedelta(hours=6)
earth2mip/networks/fcnv2_sm.py:155:15: TOR102 [*] `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
--- /home/ethan/earth2mip/earth2mip/networks/fcnv2_sm.py
+++ /home/ethan/earth2mip/earth2mip/networks/fcnv2_sm.py
@@ -150,11 +150,11 @@

     local_center = np.load(package.get("global_means.npy"))
     local_std = np.load(package.get("global_stds.npy"))

     weights_path = package.get("weights.tar")
-    weights = torch.load(weights_path, map_location=device)
+    weights = torch.load(weights_path, map_location=device, weights_only=True)
     fixed_weights = _fix_state_dict_keys(weights["model_state"], add_module=False)
     core_model.load_state_dict(fixed_weights)

     grid = earth2mip.grid.equiangular_lat_lon_grid(721, 1440)
     dt = datetime.timedelta(hours=6)
Finished checking 118 files.

Describe any alternatives you have considered

No response