(self)
| 34 | class DataLoaderTest(unittest.TestCase): |
| 35 | |
| 36 | def setUp(self): |
| 37 | super().setUp() |
| 38 | self.config = self.get_test_config(reuse_example_batch=False, per_device_batch_size=1) |
| 39 | self.config_reuse_example = self.get_test_config(reuse_example_batch=True, per_device_batch_size=1) |
| 40 | self.config_rampup = self.get_test_config( |
| 41 | reuse_example_batch=False, |
| 42 | per_device_batch_size=4.0, # This is the 'end' batch size |
| 43 | enable_rampup_batch_size=True, |
| 44 | per_device_batch_size_start=1.0, |
| 45 | per_device_batch_size_increment=1.0, |
| 46 | global_rampup_samples=60, |
| 47 | ) |
| 48 | self.mesh = Mesh(create_device_mesh(self.config), self.config.mesh_axes) |
| 49 | self.mock_data_iterator = MagicMock() |
| 50 | |
| 51 | def get_test_config(self, reuse_example_batch, **kwargs): |
| 52 | """Generate config for tests""" |
nothing calls this directly
no test coverage detected