Merge pull request #245 from jrzaurin/update-dependencies

Adjusted tests to torch 2.6 and higher

Author: jrzaurin

requirements.txt

@@ -6,7 +6,7 @@ gensim
 spacy
 opencv-contrib-python>=4.9.0.80
 tqdm
-torch >= 2.0.0, <2.6.0
+torch >= 2.0.0
 torchvision >= 0.15.0
 einops
 wrapt

tests/test_model_functioning/test_miscellaneous.py

@@ -255,7 +255,9 @@ def test_save_and_load():
     trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, batch_size=16)
     wide_weights = model.wide.wide_linear.weight.data
     trainer.save("tests/test_model_functioning/model_dir/")
-    n_model = torch.load("tests/test_model_functioning/model_dir/wd_model.pt")
+    n_model = torch.load(
+        "tests/test_model_functioning/model_dir/wd_model.pt", weights_only=False
+    )
     n_wide_weights = n_model.wide.wide_linear.weight.data
     assert torch.allclose(wide_weights, n_wide_weights)
 

tests/test_model_functioning/test_save_optimizer.py

@@ -69,7 +69,9 @@ def test_save_one_optimizer(save_state_dict):
         model_filename="model_and_optimizer.pt",
     )
 
-    checkpoint = torch.load(os.path.join(save_path, "model_and_optimizer.pt"))
+    checkpoint = torch.load(
+        os.path.join(save_path, "model_and_optimizer.pt"), weights_only=False
+    )
 
     if save_state_dict:
         new_model = WideDeep(wide=wide, deeptabular=tab_mlp)
@@ -83,7 +85,9 @@ def test_save_one_optimizer(save_state_dict):
     else:
         # This else statement is mostly testing that it runs, as it does not
         # involved loading a state_dict
-        saved_objects = torch.load(os.path.join(save_path, "model_and_optimizer.pt"))
+        saved_objects = torch.load(
+            os.path.join(save_path, "model_and_optimizer.pt"), weights_only=False
+        )
         new_model = saved_objects["model"]
         new_optimizer = saved_objects["optimizer"]
 
@@ -123,7 +127,9 @@ def test_save_multiple_optimizers(save_state_dict):
         model_filename="model_and_optimizer.pt",
     )
 
-    checkpoint = torch.load(os.path.join(save_path, "model_and_optimizer.pt"))
+    checkpoint = torch.load(
+        os.path.join(save_path, "model_and_optimizer.pt"), weights_only=False
+    )
 
     if save_state_dict:
         new_model = WideDeep(wide=wide, deeptabular=tab_mlp)
@@ -140,7 +146,9 @@ def test_save_multiple_optimizers(save_state_dict):
     else:
         # This else statement is mostly testing that it runs, as it does not
         # involved loading a state_dict
-        saved_objects = torch.load(os.path.join(save_path, "model_and_optimizer.pt"))
+        saved_objects = torch.load(
+            os.path.join(save_path, "model_and_optimizer.pt"), weights_only=False
+        )
         new_model = saved_objects["model"]
         new_optimizers = saved_objects["optimizer"]
         new_wide_opt = new_optimizers._optimizers["wide"]

tests/test_self_supervised/test_ss_miscellaneous.py

@@ -110,7 +110,9 @@ def test_save_and_load(model_type):
         save_state_dict=False,
         model_filename="ss_model.pt",
     )
-    new_model = torch.load("tests/test_self_supervised/model_dir/ss_model.pt")
+    new_model = torch.load(
+        "tests/test_self_supervised/model_dir/ss_model.pt", weights_only=False
+    )
 
     if model_type == "mlp":
         new_col_embed_module = new_model.encoder.cat_embed.embed_layers.emb_layer_col1
@@ -177,7 +179,8 @@ def test_save_model_and_optimizer(model_type, save_state_dict):
     )
 
     checkpoint = torch.load(
-        os.path.join("tests/test_self_supervised/model_dir/", "model_and_optimizer.pt")
+        os.path.join("tests/test_self_supervised/model_dir/", "model_and_optimizer.pt"),
+        weights_only=False,
     )
 
     if save_state_dict:
@@ -204,7 +207,8 @@ def test_save_model_and_optimizer(model_type, save_state_dict):
         saved_objects = torch.load(
             os.path.join(
                 "tests/test_self_supervised/model_dir/", "model_and_optimizer.pt"
-            )
+            ),
+            weights_only=False,
         )
         new_optimizer = saved_objects["optimizer"]
 
@@ -275,11 +279,15 @@ def test_save_and_load_dict(model_type):  # noqa: C901
 
     if model_type == "mlp":
         trainer2.ed_model.load_state_dict(
-            torch.load("tests/test_self_supervised/model_dir/ss_model.pt")
+            torch.load(
+                "tests/test_self_supervised/model_dir/ss_model.pt", weights_only=False
+            )
         )
     elif model_type == "transformer":
         trainer2.cd_model.load_state_dict(
-            torch.load("tests/test_self_supervised/model_dir/ss_model.pt")
+            torch.load(
+                "tests/test_self_supervised/model_dir/ss_model.pt", weights_only=False
+            )
         )
 
     if model_type == "mlp":