(self, device)
| 17 | |
| 18 | class TestContentStore(TestCase): |
| 19 | def test_basic(self, device): |
| 20 | # setup test data |
| 21 | x = torch.randn(4, device=device) |
| 22 | y = torch.randn(6, device=device) |
| 23 | z = x.view(2, 2) |
| 24 | # start writing |
| 25 | with tempfile.TemporaryDirectory() as loc: |
| 26 | writer = ContentStoreWriter(loc) |
| 27 | writer.write_tensor("x", x) |
| 28 | writer.write_tensor("y", y) |
| 29 | writer.write_tensor("z", z) |
| 30 | # do some mutation that is VC UNTRACKED |
| 31 | x.data.add_(1) |
| 32 | writer.write_tensor("x2", x) |
| 33 | writer.write_tensor("y2", y) |
| 34 | writer.write_tensor("z2", z) |
| 35 | del writer |
| 36 | |
| 37 | reader = ContentStoreReader(loc) |
| 38 | n_x = reader.read_tensor("x") |
| 39 | n_y = reader.read_tensor("y") |
| 40 | n_z = reader.read_tensor("z") |
| 41 | self.assertEqual(n_x + 1, x) |
| 42 | self.assertEqual(n_y, y) |
| 43 | self.assertEqual(n_z + 1, z) |
| 44 | self.assertEqual( |
| 45 | StorageWeakRef(n_x.untyped_storage()), |
| 46 | StorageWeakRef(n_z.untyped_storage()), |
| 47 | ) |
| 48 | n_x2 = reader.read_tensor("x2") |
| 49 | n_y2 = reader.read_tensor("y2") |
| 50 | n_z2 = reader.read_tensor("z2") |
| 51 | self.assertEqual(n_x2, x) |
| 52 | self.assertEqual(n_y2, y) |
| 53 | self.assertEqual(n_z2, z) |
| 54 | self.assertEqual( |
| 55 | StorageWeakRef(n_y2.untyped_storage()), |
| 56 | StorageWeakRef(n_y.untyped_storage()), |
| 57 | ) |
| 58 | |
| 59 | def test_scalar(self, device): |
| 60 | # Should not raise an error |
nothing calls this directly
no test coverage detected