block-ar training:   6%|████▏                                                               | 306/5000 [12:31<3:14:34,  2.49s/it, leaf=3.3571, lr=3.0e-04]Traceback (most recent call last):
step=     0 | total=11.0256 | leaf=11.0256 | ancestor=0.0000 | lr=0.00e+00 | t=3.4s
step=   100 | total=4.7289 | leaf=4.7289 | ancestor=0.0000 | lr=3.00e-04 | t=244.1s
step=   200 | total=3.5911 | leaf=3.5911 | ancestor=0.0000 | lr=2.99e-04 | t=245.7s
step=   300 | total=3.3356 | leaf=3.3356 | ancestor=0.0000 | lr=2.98e-04 | t=245.8s
  File "/home/sunhc/diffusion/sad/scripts/train_block_ar.py", line 524, in <module>
    if __name__ == "__main__":
        ^^^^^^
  File "/home/sunhc/diffusion/sad/scripts/train_block_ar.py", line 450, in main
  File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
    torch.autograd.backward(
  File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/sunhc/diffusion/sad/scripts/train_block_ar.py", line 524, in <module>
[rank0]:     if __name__ == "__main__":
[rank0]:         ^^^^^^
[rank0]:   File "/home/sunhc/diffusion/sad/scripts/train_block_ar.py", line 450, in main
[rank0]:
[rank0]:   File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/home/sunhc/miniconda3/envs/sad/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: KeyboardInterrupt