[BE] use self.assertEquals instead of str equality in test_zero1.py (#5367)

* [BE] use self.assertEquals instead of str equality in test_zero1.py * Use our own assertEqual * Remove print statements

Author: janeyx99

test/test_zero1.py

@@ -4,11 +4,12 @@
 import torch_xla.core.xla_model as xm
 from torch_xla.distributed.zero_redundancy_optimizer import ZeroRedundancyOptimizer
 from torch_xla import runtime as xr
+from torch.testing._internal.common_utils import TestCase
 
 import unittest
 
 
-class XlaZeRO1Test(unittest.TestCase):
+class XlaZeRO1Test(TestCase):
 
  @unittest.skipIf(xr.device_type() == 'TPU', "Crash on TPU")
  @unittest.skipIf(xr.device_type() == 'GPU',
@@ -33,20 +34,20 @@ def test_zero1(self):
 
  opt1.step()
  opt2.step()
- assert str(opt1.state_dict()) == str(opt2.state_dict()['base'])
+ self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
 
  s1 = opt1.state_dict()
  s2 = opt2.state_dict()
  opt1.load_state_dict(s1)
  opt2.load_state_dict(s2)
- assert str(opt1.state_dict()) == str(opt2.state_dict()['base'])
+ self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
 
  # step still runnable
  opt1.step()
  opt2.step()
  opt1.load_state_dict(s1)
  opt2.load_state_dict(s2)
- assert str(opt1.state_dict()) == str(opt2.state_dict()['base'])
+ self.assertEqual(opt1.state_dict(), opt2.state_dict()['base'])
 
  # step still runnable
  opt1.step()