fix recompute not support bug on xpu (#54367)

58fe161f · haosicheng · GitHub · 0ff446cd · 58fe161f
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

python/paddle/distributed/fleet/recompute/recompute.py python/paddle/distributed/fleet/recompute/recompute.py +2 -0

未找到文件。
--- a/python/paddle/distributed/fleet/recompute/recompute.py
+++ b/python/paddle/distributed/fleet/recompute/recompute.py
@@ -224,6 +224,8 @@ def _recompute_without_reentrant(
        cur_device = paddle.get_device()
        if 'gpu:' in cur_device:
            fw_cuda_rng_state = paddle.get_cuda_rng_state()
+        elif 'xpu:' in cur_device:
+            fw_cuda_rng_state = paddle.get_rng_state()
        elif (
            cur_device.split(':')[0]
            in paddle.device.get_all_custom_device_type()