[TEST ONLY] print statements for test_zero1.py to debug

Author: janeyx99

test/test_zero1.py

@@ -5,6 +5,7 @@
 from torch_xla.distributed.zero_redundancy_optimizer import ZeroRedundancyOptimizer
 from torch_xla import runtime as xr
 from torch.testing._internal.common_utils import TestCase
+from copy import deepcopy
 
 import unittest
 
@@ -34,19 +35,41 @@ def test_zero1(self):
 
  opt1.step()
  opt2.step()
- self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
-
  s1 = opt1.state_dict()
  s2 = opt2.state_dict()
+ print("AFTER STEPPING ONCE")
+ print("opt1.state", opt1.state)
+ print("opt1.state_dict()", s1)
+ print("opt2.state[base]", opt2.state['base'])
+ print("opt2.state_dict()[base]", s2['base'])
+ self.assertEqual(s1, s2['base'])
+
+ # s1_clone = deepcopy(s1)
+ # s2_clone = deepcopy(s2)
  opt1.load_state_dict(s1)
  opt2.load_state_dict(s2)
+ print("AFTER LOADING THE STATE_DICTs, should be same as before")
+ print("opt1.state", opt1.state)
+ print("opt1.state_dict()", opt1.state_dict())
+ print("opt2.state", opt2.state['base'])
+ print("opt2.state_dict()[base]", opt2.state_dict()['base'])
  self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
 
  # step still runnable
  opt1.step()
  opt2.step()
+ print("AFTER STEPPING AGAIN, WILL be different")
+ print("opt1.state", opt1.state)
+ print("opt1.state_dict()", opt1.state_dict())
+ print("opt2.state", opt2.state['base'])
+ print("opt2.state_dict()[base]", opt2.state_dict()['base'])
  opt1.load_state_dict(s1)
  opt2.load_state_dict(s2)
+ print("AFTER LOADING THE STATE_DICTs, should be same as before")
+ print("opt1.state", opt1.state)
+ print("opt1.state_dict()", opt1.state_dict())
+ print("opt2.state", opt2.state['base'])
+ print("opt2.state_dict()[base]", opt2.state_dict()['base'])
  self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
 
  # step still runnable

torch_patches/.torch_pin

@@ -0,0 +1 @@
+#106082