PyTorch中优化GPU上的张量报非叶张量错误如何解决

问题根因 #

你触发报错的核心原因是：对requires_grad=True的CPU张量调用.to("cuda")方法时，返回的GPU张量属于计算图的中间节点（非叶子节点），PyTorch的优化器仅支持优化叶子张量。
你之前用.detach().cpu()的方案相当于切断了梯度传播链路，自然没法更新shape的数值。

正确解决方案 #

你只需要在初始化shape张量的时候直接在GPU上创建，或者迁移后手动设置requires_grad属性即可，具体操作如下：

• 方案1：初始化shape张量时直接指定设备为cuda
  如果你是随机初始化或者从其他数据加载得到shape，直接在创建时指定device参数：

  # 示例，根据你实际的shape初始化逻辑调整
  shape = torch.randn(your_shape_size, device="cuda", requires_grad=True)
  

• 方案2：迁移到GPU后重新开启梯度
  如果你必须先在CPU上创建shape，迁移GPU后手动调用requires_grad_()标记为叶子张量：

  # 先迁移到GPU，再设置可导，此时返回的就是GPU上的叶子张量
  shape = shape.to("cuda").requires_grad_(True)
  

• 额外注意事项
  模型也同步迁移到GPU，同时真值gt_vertices也要放到同一GPU设备上，避免设备不匹配报错：

  model = model.to("cuda")
  gt_vertices = gt_vertices.to("cuda")
  

之后优化器正常传入这个处理后的shape张量即可，不需要加detach相关操作，梯度就能正常回传更新shape的数值。

内容的提问来源于stack exchange，提问作者UlucSahin