When i tried to reimplement MSELoss proposed here Loss Pytroch Implementation 4 with a real example, it gives RuntimeError: derivative for argmax is not implemented although i take argmax with nn.MSELoss and worked.
Code  def my_MSELoss(predict, true):
return ((predict - true)**2).mean()
for epoch in range(5):
losses = 0.0
for data, label in train_loader:
network.zero_grad()
predict = network(data.unsqueeze(1).float())
print(predict.shape)
#predict = predict.argmax(axis = 1)
loss = my_MSELoss(predict.argmax(axis = 1), label.float())
print(loss)
loss.backward()
optimizer.step()
loss = loss.item()
losses += loss
print(“For epoch: {}, losses are: {}”.format(epoch + 1, losses))