Repository: Lightning-Universe/lightning-transformers (language: python)

Issue:
Update imports to match with pl_bolts
## 🐛 Bug

The tests rely on an import statement for a variable name that has changed

### To Reproduce

Steps to reproduce the behavior:

1. Go to `tests/core`
2. Run `test_callbacks.py`
3. Scroll down to `==================================== ERRORS ====================================`
4. See error

```bash
==================================== ERRORS ====================================
________________ ERROR collecting tests/core/test_callbacks.py _________________
tests/core/test_callbacks.py:5: in <module>
    from pl_bolts.utils import _PL_GREATER_EQUAL_1_4_5, _TORCH_MAX_VERSION_1_8_1
E   ImportError: cannot import name '_TORCH_MAX_VERSION_1_8_1' from 'pl_bolts.utils' (/Users/runner/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pl_bolts/utils/__init__.py)
________________ ERROR collecting tests/core/test_callbacks.py _________________
ImportError while importing test module '/Users/runner/work/lightning-transformers/lightning-transformers/tests/core/test_callbacks.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
tests/core/test_callbacks.py:5: in <module>
    from pl_bolts.utils import _PL_GREATER_EQUAL_1_4_5, _TORCH_MAX_VERSION_1_8_1
E   ImportError: cannot import name '_TORCH_MAX_VERSION_1_8_1' from 'pl_bolts.utils' (/Users/runner/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pl_bolts/utils/__init__.py)
=============================== warnings summary ===============================
tests/core/test_model_step.py:7
  /Users/runner/work/lightning-transformers/lightning-transformers/tests/core/test_model_step.py:7: PytestCollectionWarning: cannot collect test class 'TestCallback' because it has a __init__ constructor (from: tests/core/test_model_step.py)
    class TestCallback(Callback):

-- Docs: https://docs.pytest.org/en/latest/warnings.html
- generated xml file: /Users/runner/work/lightning-transformers/lightning-transformers/junit/test-results-macOS-3.8-minimal.xml -
!!!!!!!!!!!!!!!!!!! Interrupted: 2 errors during collection !!!!!!!!!!!!!!!!!!!!
===================== 1 warnings, 2 error in 11.17 seconds =====================
Error: Process completed with exit code 2.
```

#### Code sample

```python
from pl_bolts.utils import _TORCH_MAX_VERSION_1_8_1
```

### Expected behavior

No error

### Environment

- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux): MacOS 10.15
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.8
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:

### Additional context

This error was found while running tests for #213 



Discussion / hints:

Fix patch:
diff --git a/lightning_transformers/utilities/imports.py b/lightning_transformers/utilities/imports.py
index 61034e2b..ec137ae5 100644
--- a/lightning_transformers/utilities/imports.py
+++ b/lightning_transformers/utilities/imports.py
@@ -17,4 +17,5 @@
 from pytorch_lightning.utilities.imports import _compare_version
 
 _BOLTS_AVAILABLE = _module_available("pl_bolts") and _compare_version("pl_bolts", operator.ge, "0.4.0")
+_BOLTS_GREATER_EQUAL_0_5_0 = _module_available("pl_bolts") and _compare_version("pl_bolts", operator.ge, "0.5.0")
 _WANDB_AVAILABLE = _module_available("wandb")


Test patch:
diff --git a/tests/core/test_callbacks.py b/tests/core/test_callbacks.py
index 2a2141a4..4bb0a39c 100644
--- a/tests/core/test_callbacks.py
+++ b/tests/core/test_callbacks.py
@@ -2,14 +2,23 @@
 import shutil
 
 import pytest
-from pl_bolts.utils import _PL_GREATER_EQUAL_1_4_5, _TORCH_MAX_VERSION_1_8_1
+from pl_bolts.utils import _PL_GREATER_EQUAL_1_4_5
 from pytorch_lightning import Trainer
 
 from lightning_transformers.core import callback
 from lightning_transformers.core.loggers import WABLogger
-from lightning_transformers.utilities.imports import _BOLTS_AVAILABLE
+from lightning_transformers.utilities.imports import _BOLTS_AVAILABLE, _BOLTS_GREATER_EQUAL_0_5_0
 from tests.core.boring_model import BoringDataModule, BoringTransformerModel
 
+if _BOLTS_GREATER_EQUAL_0_5_0:
+    from pl_bolts.utils import _TORCH_MAX_VERSION_SPARSEML as _TORCH_MAX_VERSION
+
+    MAX_VERSION = "1.10.0"
+else:
+    from pl_bolts.utils import _TORCH_MAX_VERSION_1_8_1 as _TORCH_MAX_VERSION
+
+    MAX_VERSION = "1.8.1"
+
 epoch_range_modifier = """
 - !EpochRangeModifier
    start_epoch: 0.0
@@ -56,7 +65,7 @@
 )
 @pytest.mark.skipif(not _BOLTS_AVAILABLE, reason="pytorch-lightning bolts not available")
 @pytest.mark.skipif(not _PL_GREATER_EQUAL_1_4_5, reason="pytorch-lightning version must be >= 1.4.5")
-@pytest.mark.skipif(not _TORCH_MAX_VERSION_1_8_1, reason="pytorch version must be <= 1.8.")
+@pytest.mark.skipif(not _TORCH_MAX_VERSION, reason=f"pytorch version must be <= {MAX_VERSION}.")
 def test_training_steps(max_epochs, num_processes, limit_train_batches, modifier):
     cwd = os.getcwd()
     output_dir = os.path.join(cwd, "MODELS")