paraMissing

Author: mori360

Diff for: torch/distributed/checkpoint/state_dict.py

@@ -869,6 +869,7 @@ def _device(t):
                     nonlocal device
                     if device is None:
                         device = t.device
+                    elif device != t.device:
                         raise ValueError("Device mismatch")
                 return t