I think there is a code error (Fetching the values of intermediate layers), where in two places, this: loss_history.append(loss_value) should be this: loss_history.append(loss_value.item())