Repository: Lightning-AI/pytorch-lightning
Pull request: Added getstate/setstate method for torch.save serialization
## What does this PR do?

Fixes #4114 and Fixes #4120

Still requires test, and to see if there is a less verbose way to fix serialization in the accelerator class.

Also note that this will omit any additional state variables of accelerators that inherit from this; it assumes custom accelerator will override `__getstate__` and `__setstate__`, or ensure that other objects are setup again.

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.
-->

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Unable to save model using torch.save in pytorch lightning version 0.10
## 🐛 Bug

`torch.save(trainer.model, "model.pth")` throwing error in pytorch lightning version 0.10

## Please reproduce using the following code

```
import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop. It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = pl.Trainer(max_epochs=1)
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))


torch.save(trainer.model, "model.pth")
```
`

### To Reproduce

Install pytorch lightning 0.9 and run the script mentioned above. Model will get saved successfully. 
Upgrade pytorch lightning version to 0.10 using `pip install pytorch_lightning==0.10` and run the same script, the error would be reproduced

### Expected behavior

Model should be saved as model.pth file. 

### Environment

* CUDA:
	- GPU:
	- available:         False
	- version:           10.2
* Packages:
	- numpy:             1.19.2
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0
	- pytorch-lightning: 0.10.0
	- tqdm:              4.48.2
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- ELF
	- processor:         x86_64
	- python:            3.8.2
	- version:           119~16.04.1-Ubuntu SMP Tue Sep 8 14:54:40 UTC 2020

### Additional context

Stack Trace

```
Traceback (most recent call last):
  File "encoder.py", line 43, in <module>
    torch.save(trainer.model, "model.pth")
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 364, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 466, in _save
    pickler.dump(obj)
TypeError: 'NoneType' object is not callable

```

Error while saving the model in multi gpu scenario using torch.save
## 🐛 Bug

In pytorch lightning - 0.9, while saving the model using `torch.save` in multi gpu scenario (DDP) results in error. 


```
import pytorch_lightning as pl
from sklearn.metrics import accuracy_score
from sklearn.datasets import load_iris
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split, TensorDataset


class IrisClassification(pl.LightningModule):
    def __init__(self):
        super(IrisClassification, self).__init__()
        self.fc1 = nn.Linear(4, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 3)
        self.cross_entropy_loss = nn.CrossEntropyLoss()

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.softmax(x, dim=1)
        return x

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)

        logs = {"loss": loss}
        return {"loss": loss, "log": logs}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)
        return {"val_loss": loss}

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        return {"val_loss": avg_loss}

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)
        a, y_hat = torch.max(logits, dim=1)
        test_acc = accuracy_score(y_hat.cpu(), y.cpu())
        return {"test_loss": loss, "test_acc": torch.tensor(test_acc)}

    def test_epoch_end(self, outputs):
        avg_loss = torch.stack([x["test_loss"] for x in outputs]).mean()
        avg_test_acc = torch.stack([x["test_acc"] for x in outputs]).mean()
        logs = {"test_loss": avg_loss, "test_acc": avg_test_acc}
        return {
            "avg_test_loss": avg_loss,
            "avg_test_acc": avg_test_acc,
            "log": logs,
            "progress_bar": logs,
        }


if __name__ == "__main__":
    model = IrisClassification()

    trainer = pl.Trainer(max_epochs=5, gpus=2, distributed_backend="ddp")
    #trainer = pl.Trainer(max_epochs=5)

    iris = load_iris()
    df = iris.data
    target = iris["target"]

    data = torch.Tensor(df).float()
    labels = torch.Tensor(target).long()

    data_set = TensorDataset(data, labels)
    train_set, val_set = random_split(data_set, [130, 20])
    train_set, test_set = random_split(train_set, [110, 20])
    train_loader = DataLoader(dataset=train_set, batch_size=8, num_workers=32)
    test_loader = DataLoader(dataset=test_set, batch_size=8, num_workers=32)
    val_loader = DataLoader(dataset=val_set, batch_size=8, num_workers=32)

    trainer.fit(model, train_loader, val_loader)
    trainer.test(test_dataloaders=test_loader)
    torch.save(trainer.model, "model.pth")

```

### To Reproduce

Run the script in multi gpu machine to reproduce the error. 
```
    trainer = pl.Trainer(max_epochs=5, gpus=2, distributed_backend="ddp")
```

However, running the same script in cpu works fine (model is getting saved). 

```
trainer = pl.Trainer(max_epochs=5)
```


### Expected behavior

When running with multiple gpus, model file should be dumped as `model.pth`

### Environment

* CUDA:
	- GPU:
		- Tesla K80
		- Tesla K80
		- Tesla K80
		- Tesla K80
		- Tesla K80
		- Tesla K80
		- Tesla K80
		- Tesla K80
	- available:         True
	- version:           10.2
* Packages:
	- numpy:             1.19.1
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0
	- pytorch-lightning: 0.9.0
	- tqdm:              4.50.2
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- ELF
	- processor:         x86_64
	- python:            3.8.5
	- version:           25~18.04.1-Ubuntu SMP Fri Sep 11 12:03:04 UTC 2020


### Additional context

Stack Trace:

```
Traceback (most recent call last):
  File "/home/ubuntu/ddp_test/iris.py", line 88, in <module>
    torch.save(trainer.model, "iris_gpu.pt")
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 364, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 466, in _save
    pickler.dump(obj)
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 452, in __getstate__
    self._check_default_group()
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 475, in _check_default_group
    raise RuntimeError("DDP Pickling/Unpickling are only supported "
RuntimeError: DDP Pickling/Unpickling are only supported when using DDP with the default process group. That is, when you have called init_process_group and have not passed process_group argument to DDP constructor
```

**The issue is very specific to ddp environment as the script works fine with cpu**

Tried the same classification example with pytorch 1.6 multi gpu scenario, the model is getting saved successfully.

```
import os

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from torch.nn.parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class IrisClassification(nn.Module):
    def __init__(self, dev0, dev1):
        super(IrisClassification, self).__init__()
        self.fc1 = nn.Linear(4, 10).to(dev0)
        self.fc2 = nn.Linear(10, 3).to(dev1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.softmax(x, dim=1)
        return x


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)


def demo_model_parallel(rank, world_size):
    # Dataset
    setup(rank, world_size)

    features, labels = load_iris(return_X_y=True)
    features_train, features_test, labels_train, labels_test = train_test_split(features, labels, random_state=42,
                                                                                shuffle=True)

    x_train, y_train = torch.from_numpy(features_train).float(), torch.from_numpy(labels_train).long()

    y_train = y_train.to(rank)
    x_train = x_train.to(rank)

    # create local model
    dev0 = rank * 2
    dev1 = rank * 2 + 1
    model = IrisClassification(dev0, dev1).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.CrossEntropyLoss()

    optimizer = optim.Adam(ddp_model.parameters(), lr=0.001)

    # forward pass
    outputs = ddp_model(x_train)
    # backward pass
    loss_fn(outputs, y_train).backward()
    # update parameters
    optimizer.step()

    if rank == 0:
        print("Saving Model")
        torch.save(model, "model.pth")

    cleanup()


if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    if n_gpus < 8:
        print(f"Requires at least 8 GPUs to run, but got {n_gpus}.")
    else:
        run_demo(demo_model_parallel, 2)

```

Hints from issue discussion:
```python
torch.save(trainer.model.state_dict(), "model.pth")
```
or 
```python
torch.save(trainer.get_model().state_dict(), "model.pth")
```
Yes. saving statedict works as expected in both version(0.9, 0.10). However, we are using library which dumps the entire model using `torch.save`. It was working in 0.9 and in the latest release we could find the error mentioned in the stack trace section. 
> However, we are using library which dumps the entire model using torch.save

sorry I didn't get this.

Also the code and stack trace doesn't match here.

```python
torch.save(trainer.model, "model.pth")
```
```python
Traceback (most recent call last):
  File "/tmp/test.py", line 43, in <module>
    torch.save(autoencoder, "autoencoder.pth")  <---------- here
  File "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/serialization.py", line 364, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/ubuntu/anaconda3/lib/python3.8/site-packages/torch/serialization.py", line 466, in _save
    pickler.dump(obj)
TypeError: 'NoneType' object is not callable
```
The idea is to log the entire model into mlflow using `mlflow.pytorch.log_model` - https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html . The library dumps the entire model into mlflow using `torch.save` and it is not working in 0.10 and 1.0 version. 

Updated the stack trace. 
I found two attributes that are causing pickle to throw an error:

```
trainer.model.module._results
trainer.model.module.trainer.accelerator_backend.interactive_ddp_procs
```
Deleting these keys allows the model to be saved. First one is required for single proc/GPU, both are required for multiGPU/proc (instead of interactive_ddp_procs, you'll need to delete mp_queue)

The first is an object called Result: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/core/step_result.py#L26

This doesn't override `__get_state__` `__set_state__`  functions or inherits any for serialization, thus throws an error (as it falls back to `__get_attr__`.

The second one is an accelerator class storing processes that are not cleaned up. This also happens in ddp_cpu as well. 

ddp: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp_accelerator.py#L134
ddp_cpu: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py#L56

cc @rohitgr7 @williamFalcon not entirely sure the direction we'd like to go. `trainer.save_checkpoint` via lightning handles saving, excluding these unserialisable objects. 
@SeanNaren , let's make those two picklable? 
@williamFalcon sure, will do this!
this is on 0.9. mind upgrading to 1.0?
In 1.0, https://github.com/PyTorchLightning/pytorch-lightning/issues/4114 - this error is getting reproduced in both cpu and multi gpu version. 

```
Traceback (most recent call last):                                                                                                                                                                          
  File "iris.py", line 88, in <module>
    torch.save(trainer.model, "iris_gpu.pt")
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 364, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/serialization.py", line 466, in _save
    pickler.dump(obj)
TypeError: 'NoneType' object is not callable

```


Instead of `torch.save(...)` could you replace with `trainer.save_checkpoint('model.pth')`?

https://pytorch-lightning.readthedocs.io/en/latest/weights_loading.html#manual-saving

I'll continue to investigate why torch.save doesn't work as intended, but I'm getting more verbose output:

```
Traceback (most recent call last):
  File "iris.py", line 88, in <module>
    torch.save(trainer.model, 'test.pth')
  File "/opt/conda/lib/python3.7/site-packages/torch/serialization.py", line 364, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/opt/conda/lib/python3.7/site-packages/torch/serialization.py", line 466, in _save
    pickler.dump(obj)
TypeError: can't pickle _thread.lock objects
```

This is also the same track as what @rohitgr7 is on in #4114 
@SeanNaren Thank you very much for the alternative suggestion. We are using `mlflow.pytorch` library to log the model into mlflow - https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html . Unfortunatly, the libary dumps the entire model into mlflow using `torch.save` (and there no option to save the state dict or checkpoint).

In ddp - multi gpu scenario, the same error - `TypeError: can't pickle _thread.lock objects` is getting reproduced for us.

Commit message: Added flags to ensure compatible ddp cpu environment

Patch to review:
diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 680388e812375..c3732823e3c4c 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -214,6 +214,20 @@ def init_ddp_connection(
                 torch_backend, rank=global_rank, world_size=world_size
             )
 
+    def __getstate__(self):
+        return {
+            'trainer': self.trainer,
+            'nickname': self.nickname,
+            'cluster_environment': self.cluster_environment,
+            'dist': self.dist
+        }
+
+    def __setstate__(self, d):
+        self.trainer = d['trainer']
+        self.nickname = d['nickname']
+        self.cluster_environment = d['cluster_environment']
+        self.dist = d['dist']
+
 
 # TODO: allow user to compare with string even internaly we shall use these Enum to prevent typos...
 class BackendType(Enum):
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index a45f7c0c336da..14bbf7eae316c 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -107,7 +107,7 @@ def __init__(self, *args, **kwargs):
         # optionally can be set by user
         self._example_input_array = None
         self._datamodule = None
-        self._results: Result = None
+        self._results: Optional[Result] = None
         self._current_fx_name = ''
 
     def optimizers(self):
diff --git a/pytorch_lightning/core/step_result.py b/pytorch_lightning/core/step_result.py
index 59cb8c3e6fb57..c6ed54615f7dc 100644
--- a/pytorch_lightning/core/step_result.py
+++ b/pytorch_lightning/core/step_result.py
@@ -23,6 +23,7 @@
 from pytorch_lightning.utilities.distributed import sync_ddp_if_available
 from pytorch_lightning.metrics import Metric
 
+
 class Result(Dict):
     def __init__(
         self,
@@ -89,6 +90,12 @@ def __setattr__(self, key: str, val: Union[Tensor, Any]):
 
         self[key] = val
 
+    def __getstate__(self):
+        return self
+
+    def __setstate__(self, d):
+        self.update(d)
+
     def _assert_tensor_metric(self, name: str, potential_metric: Union[bool, Tensor, None, Any]):
         if potential_metric is not None and not isinstance(potential_metric, bool):
             assert isinstance(potential_metric, Tensor), f'{name} must be a torch.Tensor'
diff --git a/tests/checkpointing/test_model_checkpoint.py b/tests/checkpointing/test_model_checkpoint.py
index 3a14f111c0dae..7f1f5f07043b4 100644
--- a/tests/checkpointing/test_model_checkpoint.py
+++ b/tests/checkpointing/test_model_checkpoint.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from distutils.version import LooseVersion
 from unittest.mock import MagicMock, Mock
 
 import yaml
@@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows",
+                    reason="Distributed training is not supported on Windows")
+@pytest.mark.skipif((platform.system() == "Darwin" and
+                     LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
+                    reason="Distributed training is not supported on MacOS before Torch 1.3.0")
+def test_model_torch_save_ddp_cpu(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer using cpu ddp."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+        distributed_backend="ddp_cpu",
+        num_processes=2,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_model_torch_save_ddp_cuda(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer using gpu ddp."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+        distributed_backend="ddp",
+        gpus=2
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)


Code review comments:
diff_hunk: @@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows",
+                    reason="Distributed training is not supported on Windows")
+@pytest.mark.skipif((platform.system() == "Darwin" and
+                     LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
path: tests/checkpointing/test_model_checkpoint.py
text: @user1:
why does it even check in env where `torch.__version__ < 1.3.0` even though lightning minimum requirements for torch is **1.3.0**?

@author:
Will remove, took it from the other ddp test here: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/tests/models/test_cpu.py#L126

diff_hunk: @@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
path: tests/checkpointing/test_model_checkpoint.py
text: @user1:
let's move all these to a special file as it is not about checkpointing and you also save the trainer
maybe tests/test_saving..py ?

diff_hunk: @@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows",
+                    reason="Distributed training is not supported on Windows")
+@pytest.mark.skipif((platform.system() == "Darwin" and
+                     LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
+                    reason="Distributed training is not supported on MacOS before Torch 1.3.0")
path: tests/checkpointing/test_model_checkpoint.py
text: @user2:
```suggestion
@user1.mark.skipif((platform.system() == "Darwin"
                     and LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
                     reason="Distributed training is not supported on MacOS before Torch 1.3.0")
```

diff_hunk: @@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows",
+                    reason="Distributed training is not supported on Windows")
+@pytest.mark.skipif((platform.system() == "Darwin" and
+                     LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
+                    reason="Distributed training is not supported on MacOS before Torch 1.3.0")
+def test_model_torch_save_ddp_cpu(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer using cpu ddp."""
+    model = EvalModelTemplate()
+    num_epochs = 4
path: tests/checkpointing/test_model_checkpoint.py
text: @user1:
maybe make it `num_epochs = 1` everywhere to reduce test run-time everywhere.

diff_hunk: @@ -512,3 +513,61 @@ def validation_epoch_end(self, outputs):
 
     # check that last one is also the best one
     assert trainer.dev_debugger.checkpoint_callback_history[-1]['epoch'] == len(monitor) - 1
+
+
+def test_model_torch_save(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows",
+                    reason="Distributed training is not supported on Windows")
+@pytest.mark.skipif((platform.system() == "Darwin" and
+                     LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
+                    reason="Distributed training is not supported on MacOS before Torch 1.3.0")
+def test_model_torch_save_ddp_cpu(tmpdir):
+    """Test to ensure torch save does not fail for model and trainer using cpu ddp."""
+    model = EvalModelTemplate()
+    num_epochs = 4
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=num_epochs,
+        distributed_backend="ddp_cpu",
+        num_processes=2,
+    )
+    temp_path = os.path.join(tmpdir, 'temp.pt')
+    trainer.fit(model)
+
+    # Ensure these do not fail
+    torch.save(trainer.model, temp_path)
+    torch.save(trainer, temp_path)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_model_torch_save_ddp_cuda(tmpdir):
path: tests/checkpointing/test_model_checkpoint.py
text: @user1:
seems like this guy is hanging...