Repository: mlflow/mlflow

Issue:
[BUG] PyTorch LightningDeprecationWarning ModelSummary and get_gpu_memory_map
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md) for additional information about bug reports. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [x] Yes. I can contribute a fix for this bug independently.
- [ ] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: No
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: 20.04
- **MLflow installed from (source or binary)**: binary
- **MLflow version (run ``mlflow --version``)**: 1.24.0
- **Python version**: 3.8.10
- **npm version, if running the dev UI**: -
- **Exact command to reproduce**: `mlflow.pytorch.autolog()`

### Describe the problem
Calling `mlflow.pytorch.autolog()` produces a warning regarding a deprecation, see logs:

### Code to reproduce issue
`mlflow.pytorch.autolog()`

### Other info / logs
```
2022/03/29 10:36:17 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: "/usr/local/lib/python3.8/dist-package
s/pytorch_lightning-1.6.0.dev0-py3.8.egg/pytorch_lightning/core/memory.py:16: LightningDeprecationWarning: `pytorch_lightning.core.memory.ge
t_memory_profile` and `pytorch_lightning.core.memory.get_gpu_memory_map` have been moved to `pytorch_lightning.utilities.memory` since v1.5 
and will be removed in v1.7."

2022/03/29 10:36:17 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: "/usr/local/lib/python3.8/dist-package
s/pytorch_lightning-1.6.0.dev0-py3.8.egg/pytorch_lightning/core/memory.py:25: LightningDeprecationWarning: `pytorch_lightning.core.memory.La
yerSummary` and `pytorch_lightning.core.memory.ModelSummary` have been moved to `pytorch_lightning.utilities.model_summary` since v1.5 and w
ill be removed in v1.7."
```


### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [x] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
Thank you for filing this @Kr4is . Please feel free to submit a backwards-compatible Version validation fix for the library deprecation and we'll be sure to take a look at the PR! Thanks!

Resolving patch:
diff --git a/mlflow/pytorch/_pytorch_autolog.py b/mlflow/pytorch/_pytorch_autolog.py
--- a/mlflow/pytorch/_pytorch_autolog.py
+++ b/mlflow/pytorch/_pytorch_autolog.py
@@ -5,7 +5,6 @@
 import shutil
 import tempfile
 import pytorch_lightning as pl
-from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.utilities import rank_zero_only
 
 from mlflow.exceptions import MlflowException
@@ -32,6 +31,10 @@
 
 
 _pl_version = Version(pl.__version__)
+if _pl_version < Version("1.5.0"):
+    from pytorch_lightning.core.memory import ModelSummary
+else:
+    from pytorch_lightning.utilities.model_summary import ModelSummary
 
 
 def _get_optimizer_name(optimizer):


Test patch: