def setUp(self): super(LearnerTest, self).setUp() # Define 4 logical CPUs on the first physical one. dist_test_utils.configure_logical_cpus() devices_gpu = tf.config.list_physical_devices('GPU') # If there are GPU devices: if devices_gpu: tf.config.experimental.set_virtual_device_configuration( devices_gpu[0], [ tf.config.LogicalDeviceConfiguration(memory_limit=10), tf.config.LogicalDeviceConfiguration(memory_limit=10), tf.config.LogicalDeviceConfiguration(memory_limit=10), tf.config.LogicalDeviceConfiguration(memory_limit=10) ])
def setUp(self): super(TrainUtilsTest, self).setUp() dist_test_utils.configure_logical_cpus()