Sorry, my bad 
Step should be called before zeroing the grad

    optimizer.zero_grad()
    optimizer.step()

import torch
import matplotlib.pyplot as plt
torch.autograd.set_detect_anomaly(True)

N=5
# load into pytorch
D = torch.rand(N, N).float()
H = torch.eye(N).float() - \
    1 / N * torch.ones(D.shape).float()
B = 0.5 * torch.matmul(torch.matmul(H, D ** 2), H)
pts = torch.rand(N, 3).float()

# declare constants
stop_loss = 1e-2
step_size = stop_loss / 3

# # emulate simulate SVD
U = torch.rand(N, N, requires_grad=True)
s = torch.diag(torch.sort(torch.rand(N), descending=True).values).requires_grad_()
Vh = torch.rand(N, N, requires_grad=True)

optimizer = torch.optim.SGD([U,s,Vh], lr=step_size)
# find embedding
embed = torch.matmul(torch.sqrt(s), Vh)
print(embed)
X_hat = embed.T[:, :3]  # select the first 3 coordinates --> x, y, z
loss_hist = []
for i in range(10):
    print(f'Iteration {i}:')
    print(f'\t U_s:{U.sum()}, s:{s.sum()}, Vh_s:{Vh.sum()}')
    # calculate loss1: how close is our SVD function?
    delta1 = (torch.matmul(U.T, U) - torch.eye(N)) + \
             (torch.matmul(Vh.T, Vh) - torch.eye(N)) + \
             (torch.matmul(torch.matmul(U, s), Vh.T) - B)

    L1 = torch.norm(delta1, p=2)

    # calculate loss2: how close is our MDS approximation?
    delta2 = X_hat - pts
    # 
    L2 = torch.norm(delta2, p=2)
    print(f'\t L1: {L1}, L2: {L2}')
    # Backprop
    loss = L1
    loss.backward()
    loss_hist.append(loss.item())
    # update

Still if I add L2 there are NaNs but guess that’s about the maths (which i didn’t check)