How to improve the performance

[Crestina2001]

I tested TAU on MMNist with the following config(the model part is copied from this repo):
hid_S:16
hid_T:256
N_S:4
N_T:4
mlp_ratio:8
drop:0
drop_path:0
spatio_kernel_enc:3
spatio_kernel_dec:3
act_inplace:true
T:10
width:64
height:64
channel:1
lr:0.001
warmup_lr:0.00001
warmup_epoch:0
epochs:200
alpha:0.1
patience:5
batch_size:16
weight_decay:0

[Crestina2001]

I trained the TAU model for 500 epochs with default hyper params, and I got:
MSE: 31.08970
MAE: 87.31856
SSIM: 0.9147
PSNR: 21.5608

It has not converged, and here is the difference visualization:

[easewkx]

I trained the TAU model for 500 epochs with default hyper params, and I got: MSE: 31.08970 MAE: 87.31856 SSIM: 0.9147 PSNR: 21.5608

It has not converged, and here is the difference visualization:

Do you have code for visualizing differences in the taxibj dataset? Can you provide it to me? Thank you.