@saeed_i You can use reduction=None parameter in the MSELoss to get what you are desiring. But I agree with @eqy , in order to take a gradient, the derivative must be with respect to a single value. While calculating the gradient you must consider the error over the complete dataset.