Are there known issues caused by ".... no current CUDA context!" warning?

Getting Started
1

I get a warning when I run any HF models wrapped by NNsight on Colab for the first time. The warning does not appear on later runs (until session restart).

with model.trace(input_ids) as trace:
    # accessing intermediate activations and final logits

"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:270.)"

I tried the below approaches but cannot get rid of the warning.

  • Creating a dummy tensor on cuda before running the model
  • Forcing PyTorch multiprocerssing to use 'spawn' instead of 'fork'
import torch.multiprocessing as mp 
if __name__ == "__main__": 
    mp.set_start_method('spawn', force=True)