torch::jit::script::Module hang on forever on loading or inference pass

username: andreabonvini
cooked: <p>Hi everyone, I’m not sure if this is the right place to post this issue or if it’s more CUDA-related.<br>
I have a super simple script which should load a <code>torch::jit::script::Module</code>  and run some forward passes on a dummy tensor.<br>
However, depending on the model I use, it gets stuck either when it tries to load the model (<code>model1</code>) or when it runs the first forward pass (<code>model2.pt</code>), and I have to manually kill the process.<br>
I’m working on a Windows 10 laptop with a NVIDIA GeForce RTX 3080 and:</p>
<ul>
<li>CUDA Toolkit 10.2.</li>
<li>LibTorch 1.10.1 for CUDA 10.2 (both Release and Debug).</li>
<li>Visual Studio 2019</li>
</ul>
<p>The output of <code>nvidia-smi</code> is:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 466.81       Driver Version: 466.81       CUDA Version: 11.3     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... WDDM  | 00000000:01:00.0 Off |                  N/A |
| N/A    0C    P5    22W /  N/A |    183MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
</code></pre>
<p>I built a VS project with <a href="https://marketplace.visualstudio.com/items?itemName=YiZhang.LibTorch001" rel="noopener nofollow ugc">this</a> template. The code is the following:</p>
<pre><code class="lang-auto">#include &lt;iostream&gt;

#include &lt;ATen/Context.h&gt;
#include &lt;torch/torch.h&gt;
#include &lt;torch/script.h&gt;

int main() {
	const torch::Device device = torch::Device(torch::kCUDA, 0);
	torch::jit::script::Module model;
	std::string MODEL_PATH = "path/to/model1.pt";
        // std::string MODEL_PATH = "path/to/model2.pt"
	// Load Model
        std::cout &lt;&lt; "Trying to load model..." &lt;&lt; std::endl;
	try {  
		model = torch::jit::load(MODEL_PATH, device);
		model.eval();
		std::cout &lt;&lt; "AI model loaded successfully." &lt;&lt; std::endl;
	}
	catch (const c10::Error&amp; e) {
		std::cerr &lt;&lt; e.what() &lt;&lt; std::endl;
	 
	}

	std::cout &lt;&lt; "Warming up model..." &lt;&lt; std::endl;

	auto dummy = torch::zeros({ 1, 3, 512, 512}).to(device);
	torch::Tensor output;
	std::vector&lt;torch::jit::IValue&gt; inputs;

	inputs.clear();
	inputs.emplace_back(dummy);


	std::cout &lt;&lt; "\tIteration: ";

	for (int i = 0; i &lt; 20; i++) {
		std::cout &lt;&lt; i &lt;&lt; ", ";
		output = model.forward(inputs).toTensor();
		torch::cuda::synchronize();
	}
}
</code></pre>
<h2>
<a name="output-for-model1-1" class="anchor" href="#output-for-model1-1"></a>Output for <code>model1</code>
</h2>
<pre><code class="lang-auto">Trying to load model...
</code></pre>
<h2>
<a name="output-for-model2-2" class="anchor" href="#output-for-model2-2"></a>Output for <code>model2</code>
</h2>
<pre><code class="lang-auto">Trying to load model...
AI model loaded successfully.
Warming up model...
        Iteration: 0,
</code></pre>
<p>I can’t figure out what could be the problem or what could I do to further investigate it.</p>

username: andreabonvini
cooked: <p><strong>Follow-up</strong>: On the RTX30 Series you MUST use at least CUDA 11.1 since it’s the first CUDA version which offers support for NVIDIA Ampere architecture (source: <a href="https://docs.nvidia.com/pdf/CUDA_Toolkit_Release_Notes.pdf" rel="noopener nofollow ugc">CUDA 11.1 Release Notes</a> chapter 1.2).<br>
After installing CUDA Toolkit 11.3 and LibTorch for CUDA 11.3 the model started working properly.</p>