i want to calculate a jacobian of an linear layer using torch.func.jacrev(). Unfortunately when doing so i get the warning shown in my minimal example below. Why does this warning occur? When I use it in my main code the runtime of different code blocks changes.
import torch
dummy = torch.tensor(1.0).cuda() # dummy to init cuda
print(torch.cuda.is_initialized())
f = torch.nn.Sequential(
torch.nn.Linear(in_features=16, out_features=16, bias=False),
).cuda()
X_k = torch.rand(16).cuda()
H = torch.func.jacrev(f)(X_k)
print(H.shape)
Output Warning:
python3.10/site-packages/torch/autograd/graph.py:769: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at ../aten/src/ATen/cuda/CublasHandlePool.cpp:135.)
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
@ptrblck, could you clarify how the warning relates to multi-threading?
If I understand correctly, the warning appears even when running the “vanilla” code with a single GPU and a single thread. Do you have any suggestions on adjustments to prevent the warning from appearing?
PyTorch uses a separate thread in the backward pass. If no other CUDA operation was used before the cuBLAS call, the warning will be raised as we will set the primary context to the current one.
Hi, I encountered a similar problem when I was running the WGAN-GP code. I have already set the primary context to the current one for multi-threaded use cases, but the problem still exists.(There is no error when I remove the gradient penalty part in the program.)The following is the content of the error:D:\anaconda3\envs\py312\Lib\site-packages\torch\autograd\graph.py:825: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context… (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\cuda\CublasHandlePool.cpp:135.)
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass Could you give me some advice, thanks!!!