Fix dependencies

Signed-off-by: Beat Buesser <[email protected]>

Author: beat-buesser

tests/attacks/evasion/test_overload_attack.py

@@ -31,9 +31,11 @@
 def test_generate(art_warning):
     try:
         import torch
+        import ultralytics
         from ultralytics import YOLO
 
         torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
+        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
 
         model = YOLO("yolov5su.pt")
         py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)
@@ -67,9 +69,11 @@ def test_generate(art_warning):
 def test_check_params(art_warning):
     try:
         import torch
+        import ultralytics
         from ultralytics import YOLO
 
         torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
+        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
 
         model = YOLO("yolov5su.pt")
         py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True)

tests/attacks/evasion/test_steal_now_attack_later.py

@@ -34,9 +34,11 @@ def test_generate(art_warning):
         # To avoid an import error with the TF 1.x pipeline, it is imported only within the function scope.
         import torch
         import requests
+        import ultralytics
         from ultralytics import YOLO
 
         torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
+        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
 
         model = YOLO("yolov8m")
         py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)
@@ -199,9 +201,11 @@ def test_check_params(art_warning):
         # To avoid an import error with the TF 1.x pipeline, it is imported only within the function scope.
         import torch
         import requests
+        import ultralytics
         from ultralytics import YOLO
 
         torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])
+        torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
 
         model = YOLO("yolov8m")
         py_model = PyTorchYolo(model=model, input_shape=(3, 640, 640), channels_first=True, is_yolov8=True)