自定义损失函数在反向传播时变为零

import torch
import torch.nn as nn
import torch.optim as optim

# 自定义损失函数
class CustomLoss(nn.Module):
    def __init__(self):
        super(CustomLoss, self).__init__()

    def forward(self, input, target):
        return torch.mean((input - target) ** 2)

# 创建模型和优化器
model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = CustomLoss()

# 假设有一些输入数据和目标数据
inputs = torch.randn(32, 10)
targets = torch.randn(32, 1)

# 前向传播
outputs = model(inputs)
loss = criterion(outputs, targets)

# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 检查梯度是否为零
for param in model.parameters():
    if param.grad is not None:
        print(param.grad.norm())  # 打印梯度的范数，检查是否接近零