fix: Fix error introduced by merge

CAREamics · Jul 3, 2024 · d0369f7 · d0369f7
1 parent 67ace1c
commit d0369f7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/lightning/test_lightning_api.py b/tests/lightning/test_lightning_api.py
@@ -130,4 +130,4 @@ def test_smoke_n2v_2d_tiling(tmp_path, minimum_configuration):
     # predict
     predicted = trainer.predict(model, datamodule=predict_data)
     predicted_stitched = convert_outputs(predicted, tiled=True)
-    assert predicted_stitched.shape == val_array.shape
+    assert predicted_stitched[0].shape[-2:] == val_array.shape