Update docstrings

Signed-off-by: Beat Buesser <[email protected]>
Trusted-AI · Dec 23, 2023 · 4abbedb · 4abbedb
1 parent 7a8cce2
commit 4abbedb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/art/estimators/object_detection/pytorch_object_detector.py b/art/estimators/object_detection/pytorch_object_detector.py
@@ -298,8 +298,8 @@ def _get_losses(
         """
         self._model.train()
 
-        self.set_dropout(False)
-        self.set_multihead_attention(False)
+        self.set_dropout(train=False)
+        self.set_multihead_attention(train=False)
 
         # Apply preprocessing and convert to tensors
         x_preprocessed, y_preprocessed = self._preprocess_and_convert_inputs(x=x, y=y, fit=False, no_grad=False)

diff --git a/tests/estimators/object_detection/test_pytorch_detection_transformer.py b/tests/estimators/object_detection/test_pytorch_detection_transformer.py
@@ -139,7 +139,7 @@ def test_loss_gradient(art_warning, get_pytorch_detr):
         )
 
         print("expected_gradients1")
-        print(expected_gradients1)
+        print(grads[0, 0, 10, :32])
 
         np.testing.assert_array_almost_equal(grads[0, 0, 10, :32], expected_gradients1, decimal=1)
 
@@ -181,7 +181,7 @@ def test_loss_gradient(art_warning, get_pytorch_detr):
         )
 
         print("expected_gradients2")
-        print(expected_gradients2)
+        print(grads[1, 0, 10, :32])
 
         np.testing.assert_array_almost_equal(grads[1, 0, 10, :32], expected_gradients2, decimal=2)