Fixing format for pytorch yolo wrapper test.

Signed-off-by: Kieran Fraser <[email protected]>

Author: kieranfraser

tests/estimators/object_detection/test_pytorch_yolo_loss_wrapper.py

@@ -46,15 +46,18 @@ def test_yolov8_loss_wrapper():
         boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]])  # [x1, y1, x2, y2]
         labels = torch.zeros(2, dtype=torch.long)  # Use class 0 for testing
         targets.append({"boxes": boxes, "labels": labels})"""
-    targets = torch.tensor([[ 0.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582],
-        [ 0.0000, 20.0000,  0.2487,  0.4062,  0.4966,  0.5787],
-        [ 0.0000, 20.0000,  0.5667,  0.2772,  0.0791,  0.2313],
-        [ 0.0000, 20.0000,  0.1009,  0.1955,  0.2002,  0.0835],
-        [ 1.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582],
-        [ 1.0000, 20.0000,  0.2487,  0.4062,  0.4966,  0.5787],
-        [ 1.0000, 20.0000,  0.5667,  0.2772,  0.0791,  0.2313],
-        [ 1.0000, 20.0000,  0.1009,  0.1955,  0.2002,  0.0835]])
-
+    targets = torch.tensor(
+        [
+            [0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [0.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [0.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [0.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+            [1.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [1.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [1.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [1.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+        ]
+    )
 
     # Test training mode
     losses = wrapper(x, targets)
@@ -108,14 +111,18 @@ def test_yolov10_loss_wrapper():
         boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]])  # [x1, y1, x2, y2]
         labels = torch.zeros(2, dtype=torch.long)  # Use class 0 for testing
         targets.append({"boxes": boxes, "labels": labels})"""
-    targets = torch.tensor([[ 0.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582],
-        [ 0.0000, 20.0000,  0.2487,  0.4062,  0.4966,  0.5787],
-        [ 0.0000, 20.0000,  0.5667,  0.2772,  0.0791,  0.2313],
-        [ 0.0000, 20.0000,  0.1009,  0.1955,  0.2002,  0.0835],
-        [ 1.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582],
-        [ 1.0000, 20.0000,  0.2487,  0.4062,  0.4966,  0.5787],
-        [ 1.0000, 20.0000,  0.5667,  0.2772,  0.0791,  0.2313],
-        [ 1.0000, 20.0000,  0.1009,  0.1955,  0.2002,  0.0835]])
+    targets = torch.tensor(
+        [
+            [0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [0.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [0.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [0.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+            [1.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582],
+            [1.0000, 20.0000, 0.2487, 0.4062, 0.4966, 0.5787],
+            [1.0000, 20.0000, 0.5667, 0.2772, 0.0791, 0.2313],
+            [1.0000, 20.0000, 0.1009, 0.1955, 0.2002, 0.0835],
+        ]
+    )
 
     # Test training mode
     losses = wrapper(x, targets)
@@ -236,7 +243,7 @@ def loss(self, items):
     wrapper.train()
     # Dummy input and targets
     x = torch.zeros((1, 3, 416, 416))
-    targets = torch.tensor([[ 0.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582]])
+    targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]])
     losses = wrapper(x, targets)
     assert set(losses.keys()) == {"loss_total", "loss_box", "loss_cls", "loss_dfl"}
     assert losses["loss_total"].item() == 6.0  # sum([1.0, 2.0, 3.0])
@@ -281,7 +288,7 @@ def loss(self, items):
     wrapper = PyTorchYoloLossWrapper(test_model, name="yolov8n")
     wrapper.train()
     x = torch.zeros((1, 3, 416, 416))
-    targets = torch.tensor([[ 0.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582]])
+    targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]])
     losses = wrapper(x, targets)
     assert set(losses.keys()) == {"loss_total", "loss_box", "loss_cls", "loss_dfl"}
     assert losses["loss_total"].item() == 6.0
@@ -456,7 +463,7 @@ def loss(self, items):
     for batch_size in batch_sizes:
         for box_count in box_counts:
             x = torch.zeros((batch_size, 3, 416, 416))
-            targets = torch.tensor([[ 0.0000, 20.0000,  0.7738,  0.3919,  0.4525,  0.7582]]*batch_size)
+            targets = torch.tensor([[0.0000, 20.0000, 0.7738, 0.3919, 0.4525, 0.7582]] * batch_size)
             losses = wrapper(x, targets)
 
             # Verify loss structure