Update imports to match with pl_bolts

[mathemusician]

🐛 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

==================================== 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

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