update

Author: justusschock

tests/tests_pytorch/loggers/test_all.py

@@ -287,7 +287,7 @@ def _test_logger_initialization(tmp_path, logger_class):
 
 
 @mock.patch.dict(os.environ, {})
-@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
+@mock.patch("pytorch_lightning_enterprise.loggers.mlflow._get_resolve_tags", Mock())
 def test_logger_with_prefix_all(mlflow_mock, wandb_mock, comet_mock, neptune_mock, monkeypatch, tmp_path):
     """Test that prefix is added at the beginning of the metric keys."""
     prefix = "tmp"

tests/tests_pytorch/models/test_tpu.py

@@ -235,7 +235,7 @@ def test_tpu_misconfiguration(devices, tpu_available):
 
 
 @pytest.mark.skipif(XLAAccelerator.is_available(), reason="test requires missing TPU")
-@mock.patch("lightning.fabric.accelerators.xla._using_pjrt", return_value=True)
+@mock.patch("pytorch_lightning_enterprise.accelerators.xla._using_pjrt", return_value=True)
 def test_exception_when_no_tpu_found(_, xla_available):
     """Test if exception is thrown when xla devices are not available."""
     with pytest.raises(MisconfigurationException, match="XLAAccelerator` can not run on your system"):