def test_auto_device_count(_): assert CPUAccelerator.auto_device_count() == 1 assert CUDAAccelerator.auto_device_count() == 2 assert TPUAccelerator.auto_device_count() == 8 assert IPUAccelerator.auto_device_count() == 4
def test_broadcast(rank): trainer = Trainer(tpu_cores=8) backend = TPUAccelerator(trainer) obj = ("ver_0.5", "logger_name", rank) result = backend.broadcast(obj) assert result == ("ver_0.5", "logger_name", 0)
def test_auto_device_count(device_count_mock): assert CPUAccelerator.auto_device_count() == 1 assert GPUAccelerator.auto_device_count() == 2 assert TPUAccelerator.auto_device_count() == 8 assert IPUAccelerator.auto_device_count() == 4