Fix failing Github tests (#726)

Summary:
Pull Request resolved: #726

(1) CIFAR10 cuda tests were failing with the following message. [Example](https://github.com/pytorch/opacus/actions/runs/13079011269/job/36498002040#step:7:757).

> _pickle.UnpicklingError: Weights only load failed [...] In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.

Specified weights_only=False in torch.load() and tests are passing.

(2) A similar issue was happening with multiple unit tests due to the change in behavior of `weights_only`. [Example](https://github.com/pytorch/opacus/actions/runs/13080274224/job/36502044189). Changed the `weights_only` setting where necessary.

(3) Black was requesting addition of commas when listing mutliple variable types. (Locally there were no issues, appeared only GithubActions). [Example](https://github.com/pytorch/opacus/actions/runs/13080274224/job/36502043210).

Reviewed By: EnayatUllah

Differential Revision: D68973109

fbshipit-source-id: 7fbe9702dcfc72b40fe95bb8a749d2163525231b

Author: iden-kalemaj

.github/workflows/ci_gpu.yml

@@ -84,9 +84,9 @@ jobs:
           mkdir -p runs/cifar10/test-reports
           pip install tensorboard
           python examples/cifar10.py --lr 0.1 --sigma 1.5 -c 10 --batch-size 2000 --epochs 10 --data-root runs/cifar10/data --log-dir runs/cifar10/logs --device cuda
-          python -c "import torch; model = torch.load('model_best.pth.tar'); exit(0) if (model['best_acc1']>0.4 and model['best_acc1']<0.49) else exit(1)"
+          python -c "import torch; model = torch.load('model_best.pth.tar', weights_only=False); exit(0) if (model['best_acc1']>0.4 and model['best_acc1']<0.49) else exit(1)"
           python examples/cifar10.py --lr 0.1 --sigma 1.5 -c 10 --batch-size 2000 --epochs 10 --data-root runs/cifar10/data --log-dir runs/cifar10/logs --device cuda --grad_sample_mode no_op
-          python -c "import torch; model = torch.load('model_best.pth.tar'); exit(0) if (model['best_acc1']>0.4 and model['best_acc1']<0.49) else exit(1)"
+          python -c "import torch; model = torch.load('model_best.pth.tar', weights_only=False); exit(0) if (model['best_acc1']>0.4 and model['best_acc1']<0.49) else exit(1)"
 
       - name: Store CIFAR10 test results
         uses: actions/upload-artifact@v4

benchmarks/tests/test_utils.py

@@ -194,7 +194,7 @@ def pickle_data_and_config(
     ],
 )
 def test_save_results(
-    pickle_data_and_config: Tuple[Dict[str, Any], Dict[str, Any]]
+    pickle_data_and_config: Tuple[Dict[str, Any], Dict[str, Any]],
 ) -> None:
     """Tests saving benchmark results.
 

opacus/grad_sample/utils.py

@@ -27,7 +27,7 @@
 
 
 def register_grad_sampler(
-    target_class_or_classes: Union[Type[nn.Module], Sequence[Type[nn.Module]]]
+    target_class_or_classes: Union[Type[nn.Module], Sequence[Type[nn.Module]]],
 ):
     """
     Registers the decorated function as the ``grad_sampler`` of ``target_class_or_classes``, which is
@@ -56,7 +56,7 @@ def decorator(f):
 
 
 def register_norm_sampler(
-    target_class_or_classes: Union[Type[nn.Module], Sequence[Type[nn.Module]]]
+    target_class_or_classes: Union[Type[nn.Module], Sequence[Type[nn.Module]]],
 ):
     """
     Registers the decorated function as the ``norm_sampler`` of ``target_class_or_classes``, which is

opacus/privacy_engine.py

@@ -603,7 +603,7 @@ def load_checkpoint(
         module_load_dict_kwargs: Optional[Dict[str, Any]] = None,
         torch_load_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Dict:
-        checkpoint = torch.load(path, **(torch_load_kwargs or {}))
+        checkpoint = torch.load(path, **(torch_load_kwargs or {}), weights_only=False)
         module.load_state_dict(
             checkpoint["module_state_dict"], **(module_load_dict_kwargs or {})
         )

opacus/utils/per_sample_gradients_utils.py

@@ -39,7 +39,7 @@ def clone_module(module: nn.Module) -> nn.Module:
     with io.BytesIO() as bytesio:
         torch.save(module, bytesio)
         bytesio.seek(0)
-        module_copy = torch.load(bytesio)
+        module_copy = torch.load(bytesio, weights_only=False)
     return module_copy