We could remove the dependency on batch size by dividing the loss by batch size. Would that work?
criterion = nn.MSELoss(reduction='sum')
out = model(x)
loss = criterion(out, y)
loss = loss / x.shape[0]