context.context().set_logical_device_configuration(physical_gpus[1], [
        context.LogicalDeviceConfiguration(memory_limit=1024),
        context.LogicalDeviceConfiguration(memory_limit=1024)