.

Author: Borda

pytorch_lightning/overrides/data_parallel.py

@@ -285,7 +285,7 @@ def _worker(i, module, input, kwargs, device=None):
                 if output is None:
                     warn_missing_output(fx_called)
 
-                if output is not None and module.distrib_type in ("dp", "ddp2"):
+                if output is not None and module._distrib_type in ('dp', 'ddp2'):
                     auto_squeeze_dim_zeros(output)
                 # ---------------
 

pytorch_lightning/utilities/enums.py

@@ -50,7 +50,7 @@ class DistributedType(LightningEnum):
     >>> DistributedType.DDP == 'ddp'
     True
     >>> # which is case invariant
-    >>> DistributedType.DDP2 == 'DDP2'
+    >>> DistributedType.DDP2 in ('ddp2', )
     True
     """
     DP = 'dp'
@@ -69,7 +69,7 @@ class DeviceType(LightningEnum):
     >>> DeviceType.GPU == 'GPU'
     True
     >>> # which is case invariant
-    >>> DeviceType.TPU == 'tpu'
+    >>> DeviceType.TPU in ('tpu', 'CPU')
     True
     """
     CPU = 'CPU'