device_guard
可以切换当前的 device 为输入指定的 device。
备注
该 API 目前仅支持动态图模式。
参数
device (PlaceLike) - 指定的 device,可以是形如 "cpu", "gpu:0" 之类的设备描述字符串,也可以是 paddle.CUDAPlace(0), paddle.CPUPlace() 之类的设备实例。
代码示例
>>> import paddle
>>> # Set the global default device to CPU
>>> paddle.set_device("cpu")
>>> # Temporarily switch to GPU:0 using device_guard with string input
>>> with paddle.device.device_guard("gpu:0"):
... x = paddle.randn([4, 4]) # Create a Tensor on GPU:0
... x = x.tanh() * 2 # Perform computation on GPU:0
... print(x.place) # Check the device of the Tensor
Place(gpu:0)
>>> # Set the global default device to GPU:0
>>> paddle.set_device("gpu:0")
>>> # Temporarily switch to CPU using device_guard with Place object (CPUPlace)
>>> cpu_place = paddle.CPUPlace()
>>> with paddle.device.device_guard(cpu_place):
... x = paddle.randn([4, 4]) # Create a Tensor on CPU
... x = x.tanh() * 2 # Perform computation on CPU
... print(x.place)
Place(cpu)