RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu!

[KDD2018]

deepspeed启用zero3时，在多个迭代训练之后出现RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu! zero3参数配置如下：

"zero_optimization": { "stage": 3, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "offload_param": { "device": "cpu", "pin_memory": true }, "overlap_comm": true, "contiguous_gradients": true, "sub_group_size": 1e9, "reduce_bucket_size": "auto", "stage3_prefetch_bucket_size": "auto", "stage3_param_persistence_threshold": "auto", "stage3_max_live_parameters": 1e9, "stage3_max_reuse_distance": 1e9, "stage3_gather_16bit_weights_on_model_save": true },