TorchScript model doesn't work with autocast

Shirish_Vogga: <p>Hi I want to use autocast with a script model and had the following error.</p>
<p>import torch<br>
import torch.nn as nn<br>
from torch.amp import autocast</p>
<p>class SimpleCNN(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(SimpleCNN, self).<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)</p>
<pre><code>def forward(self, x):
    with autocast(device_type="cuda", enabled=True):
        return self.conv1(x)
</code></pre>
<p>device = torch.device(‘cuda’)</p>
<h1><a name="p-459470-create-an-instance-of-the-network-1" class="anchor" href="#p-459470-create-an-instance-of-the-network-1"></a>Create an instance of the network</h1>
<p>net = SimpleCNN()<br>
net.to(device)</p>
<h1><a name="p-459470-create-a-sample-input-tensor-2" class="anchor" href="#p-459470-create-a-sample-input-tensor-2"></a>Create a sample input tensor</h1>
<p>input_tensor = torch.randn(1, 3, 28, 28)<br>
input_tensor = input_tensor.to(device)</p>
<h1><a name="p-459470-pass-the-input-through-the-network-3" class="anchor" href="#p-459470-pass-the-input-through-the-network-3"></a>Pass the input through the network</h1>
<p>output = net(input_tensor)<br>
print(output.shape)</p>
<h1><a name="p-459470-pass-the-input-through-the-script-network-4" class="anchor" href="#p-459470-pass-the-input-through-the-script-network-4"></a>Pass the input through the script network</h1>
<p>script_net = torch.jit.script(net)<br>
output2 = script_net(input_tensor)<br>
print(output2.shape)</p>
<p>the error thrown out is as shown below</p>
<pre><code class="lang-auto">torch.Size([1, 6, 24, 24])
Traceback (most recent call last):
  File "D:\workspace_tf\Einstein_reg_8\src\neosoft\misc\SimpleCNN.py", line 39, in &lt;module&gt;
    output2 = script_net(input_tensor)
  File "C:\Python39\lib\site-packages\torch\nn\modules\module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Python39\lib\site-packages\torch\nn\modules\module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "D:\workspace_tf\Einstein_reg_8\src\neosoft\misc\SimpleCNN.py", line 21, in forward
            # x = torch.relu(self.fc1(x))
            # return x
            return self.conv1(x)
                   ~~~~~~~~~~ &lt;--- HERE
  File "C:\Python39\lib\site-packages\torch\nn\modules\conv.py", line 554, in forward
    def forward(self, input: Tensor) -&gt; Tensor:
        return self._conv_forward(input, self.weight, self.bias)
               ~~~~~~~~~~~~~~~~~~ &lt;--- HERE
  File "C:\Python39\lib\site-packages\torch\nn\modules\conv.py", line 549, in _conv_forward
                self.groups,
            )
        return F.conv2d(
               ~~~~~~~~ &lt;--- HERE
            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
        )
RuntimeError: Input type (struct c10::Half) and bias type (float) should be the same
</code></pre>

ptrblck: <p>Try to wrap <code>autocast</code> around the forward pass as shown in the <a href="https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#adding-torch-autocast">docs</a> and your model should work. <code>TorchScript</code> is also in maintenance mode and we recommend using <code>torch.compile</code> instead.</p>

ktsumura: <p>Thank you for the reply. I would like to use autocast partially, as my model uses PCA, which doesn’t support half precision. I will try torch.compile.</p>