TypeError: unsupported operand type(s) for *: 'NoneType' and 'float'

Melani_Cook: <p>I have a relatively large architecture. Here is the training section where I use <code>torch.cuda.amp</code> for my model:</p>
<blockquote>
<pre><code>def train(epoch):
    # model in training mode
    modelstate.model.train()
    # initialization
    total_loss = 0
    total_batches = 0
    total_points = 0
    if torch.cuda.is_available():
        
        scaler = torch.cuda.amp.GradScaler()

    for i, (u, y) in enumerate(loader_train):
        u = u.to(device)#torch.Size([B, D_in, T])
        y = y.to(device)
        
        # set the optimizer
                    # set the optimizer
        modelstate.optimizer.zero_grad()
        if torch.cuda.is_available():
            with torch.autocast(device_type='cuda', dtype=torch.float32) and torch.backends.cudnn.flags(enabled=False):
                loss_ = modelstate.model(u, y)
            diff_params = [p for p in modelstate.model.m.parameters() if p.requires_grad]    
            scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss_),
                                                    inputs=diff_params,
                                                    create_graph=True,
                                                    retain_graph=True,
                                                    allow_unused=True #Whether to allow differentiation of unused parameters.
                                                    )
             
            
            #find NaN values in the architecture
            for submodule in modelstate.model.m.modules():
                submodule.register_forward_hook(nan_hook)
                
            inv_scale = 1./scaler.get_scale()
            
            #grad_params = [ p * inv_scale if p is not None and not torch.isnan(p).any() else torch.tensor(0, device=device, dtype=torch.float32) for p in scaled_grad_params ]
            grad_params = [p * inv_scale for p in scaled_grad_params]
            with torch.autocast(device_type='cuda', dtype=torch.float32):
                #grad_norm = torch.tensor(0, device=grad_params[0].device, dtype=grad_params[0].dtype)
                grad_norm = 0
                for grad in grad_params:
                    grad_norm += grad.pow(2).sum()
                    grad_norm = grad_norm**0.5
                # Compute the L2 Norm as penalty and add that to loss
                loss_ = loss_ + grad_norm
</code></pre>
</blockquote>
<p>Here is the error message :</p>
<blockquote>
<p>/tmp/ipykernel_100588/1815332848.py in (.0)<br>
3006<br>
3007                 <span class="hashtag-raw">#grad_params</span> = [ p * inv_scale if p is not None and not torch.isnan(p).any() else torch.tensor(0, device=device, dtype=torch.float32) for p in scaled_grad_params ]<br>
 → 3008                 grad_params = [p * inv_scale for p in scaled_grad_params]<br>
3009                 with torch.autocast(device_type=‘cuda’, dtype=torch.float32):<br>
3010                     <span class="hashtag-raw">#grad_norm</span> = torch.tensor(0, device=grad_params[0].device, dtype=grad_params[0].dtype)</p>
<p>TypeError: unsupported operand type(s) for *: ‘NoneType’ and ‘float’</p>
</blockquote>
<p>Before, I changed the problematic line with this <code>grad_params = [ p * inv_scale if p is not None and not torch.isnan(p).any() else torch.tensor(0, device=device, dtype=torch.float32) for p in scaled_grad_params ]</code> and I wasn’t getting this error.</p>
<p>However, I wasn’t happy with the results of my model and I read in a debugging blog that I must check the gradient for every parameter and if they come up empty that is where I must take a closer look. The blog didn’t go further why it is a problem. I am wondering how to deal with this tyep of error in general?</p>

ptrblck: <p>The issue is raised in the mentioned line of code since you are using <code>allow_unused=True</code> in the <code>grad</code> call, which will return <code>None</code> gradients for unused parameters as seen here:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)
        self.unused = nn.Linear(1, 1)
        
    def forward(self, x):
        x = self.fc1(x)
        return x
    
model = MyModel()
x = torch.randn(1, 1)
out = model(x)
loss = out.mean()

grads = torch.autograd.grad(outputs=loss,
                            inputs=model.parameters(),
                            create_graph=True,
                            retain_graph=True,
                            allow_unused=True
                            )
print(grads)
# (tensor([[0.2469]]), tensor([1.]), None, None)
</code></pre>
<p>You might thus need to check for <code>None</code> gradients before trying to unscale them.</p>