/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
t2sel=40   cos=0.9145 0.8368 0.9153  mean=0.8889
t2sel=80   cos=0.9631 0.8890 0.9561  mean=0.9361
t2sel=160  cos=0.9651 0.8781 0.8942  mean=0.9124
t2sel=320  cos=0.9651 0.8781 0.8748  mean=0.9060
T2_PROBE_DONE