Such a good work and thanks for the code!
But I met a problem in unet step when run python train_svd.py

Problem1:

So, I select the solution of inp_noisy_latents = inp_noisy_latents.unsqueeze(0) before unet-pred step and let the 5d to 6d.
but probblem2 occured and is conflicting with problem1.

could u pls help me and thx in advance.