test: enhance tensorboard log graph

Author: KAVYANSHTYAGI

tests/tests_fabric/loggers/test_tensorboard.py

@@ -146,30 +146,47 @@ def test_tensorboard_log_hparams_and_metrics(tmp_path):
     logger.log_hyperparams(hparams, metrics)
 
 
+@pytest.mark.parametrize("model_cls", [BoringModel, pytest.importorskip("lightning.pytorch.demos.boring_classes").BoringModel])
 @pytest.mark.parametrize("example_input_array", [None, torch.rand(2, 32)])
-def test_tensorboard_log_graph(tmp_path, example_input_array):
+def test_tensorboard_log_graph(tmp_path, example_input_array, model_cls):
     """Test that log graph works with both model.example_input_array and if array is passed externally."""
-    # TODO(fabric): Test both nn.Module and LightningModule
-    # TODO(fabric): Assert _apply_batch_transfer_handler is calling the batch transfer hooks
-    model = BoringModel()
+    model = model_cls()
     if example_input_array is not None:
         model.example_input_array = None
 
     logger = TensorBoardLogger(tmp_path)
     logger._experiment = Mock()
-    logger.log_graph(model, example_input_array)
-    if example_input_array is not None:
-        logger.experiment.add_graph.assert_called_with(model, example_input_array)
-    logger._experiment.reset_mock()
 
-    # model wrapped in `FabricModule`
-    wrapped = _FabricModule(model, strategy=Mock())
-    logger.log_graph(wrapped, example_input_array)
-    if example_input_array is not None:
-        logger.experiment.add_graph.assert_called_with(model, example_input_array)
+    if isinstance(model, torch.nn.Module) and hasattr(model, "_apply_batch_transfer_handler"):
+        with (
+            mock.patch.object(model, "_on_before_batch_transfer", return_value=example_input_array) as before_mock,
+            mock.patch.object(model, "_apply_batch_transfer_handler", return_value=example_input_array) as transfer_mock,
+        ):
+            logger.log_graph(model, example_input_array)
+            logger._experiment.reset_mock()
+            wrapped = _FabricModule(model, strategy=Mock())
+            logger.log_graph(wrapped, example_input_array)
+            if example_input_array is not None:
+                assert before_mock.call_count == 2
+                assert transfer_mock.call_count == 2
+                logger.experiment.add_graph.assert_called_with(model, example_input_array)
+            else:
+                before_mock.assert_not_called()
+                transfer_mock.assert_not_called()
+                logger.experiment.add_graph.assert_not_called()
+    else:
+        logger.log_graph(model, example_input_array)
+        if example_input_array is not None:
+            logger.experiment.add_graph.assert_called_with(model, example_input_array)
+        logger._experiment.reset_mock()
+
+        wrapped = _FabricModule(model, strategy=Mock())
+        logger.log_graph(wrapped, example_input_array)
+        if example_input_array is not None:
+            logger.experiment.add_graph.assert_called_with(model, example_input_array)
 
 
-@pytest.mark.skipif(not _TENSORBOARD_AVAILABLE, reason=str(_TENSORBOARD_AVAILABLE))
+@pytest.mark.skipif(not _TENSORBOARD_AVAILABLE, reason="tensorboard is required")
 def test_tensorboard_log_graph_warning_no_example_input_array(tmp_path):
     """Test that log graph throws warning if model.example_input_array is None."""
     model = BoringModel()