shape mismatch indexError in contiformer example

[ThpaperUp]

Hello, thank you for releasing the code for your work. I got the following issue when running your example using python spiral.py --adjoint=1 --visualize=1 --niters=10000 --model_name Contiformer --noise_a=0.02 --cc=2 --train_dir ./spiral_contiformer, could you please help to take a look? Thank you very much

2023-12-22 18:23:11,075 ./spiral_contiformer/log INFO Loaded ckpt from ./spiral_contiformer/ckpt_Contiformer.pth
2023-12-22 18:23:25,199 ./spiral_contiformer/log INFO Iter: 4, running loss: 5050.6167
2023-12-22 18:23:25,219 ./spiral_contiformer/log INFO Stored ckpt at ./spiral_contiformer/ckpt_Contiformer.pth
Traceback (most recent call last):
File "/content/drive/MyDrive/code/PhysioPro/spiral.py", line 470, in
pred_x = model(samp_trajs, orig_ts, idx=test_idx)[0]
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/code/PhysioPro/spiral.py", line 329, in forward
out, _ = self.encoder(input, orig_ts.unsqueeze(0).repeat(self.batch_size, 1).float(),
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/code/PhysioPro/physiopro/network/contiformer.py", line 236, in forward
enc_output, enc_slf_attn = self.slf_attn(x, x, x, t, mask=mask)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/code/PhysioPro/physiopro/network/contiformer.py", line 55, in forward
q = self.w_qs(q, t).view(sz_b, len_q, len_q, -1, n_head, d_k)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/code/PhysioPro/physiopro/module/linear.py", line 90, in forward
coeffs, intervals = self.pre_integrals(x, t)
File "/content/drive/MyDrive/code/PhysioPro/physiopro/module/linear.py", line 67, in pre_integrals
xx[r.numpy(), sort_index.reshape(-1).cpu().numpy(), :] = x.reshape(x.size(0) * x.size(1), -1) # [B, TT, D]
IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [15000], [9600]

[chenyuqi990215]

It seems that the code has run for 4 iterations. Maybe you can remove all the saved temporary files and re-run the code. Feel free to contact me if the problem still exists.

[yj5049]

I got the same issue when running spiral.py --adjoint=1 --visualize=1 --niters=10000 --model_name Contiformer --noise_a=0.02 --cc=2 --train_dir ./spiral_contiformer

[chenyuqi990215]

Sorry for my mistake and we have updated the code.