[PyTorch] Debug weight matrix usages for dgrad GEMM (#1637)

Make sure that weight matrix has required usages for dgrad GEMM
Signed-off-by: Tim Moon <tmoon@nvidia.com>

transformer_engine/pytorch/module/layernorm_linear.py

@@ -327,9 +327,8 @@ class _LayerNormLinear(torch.autograd.Function):
                         ln_out.update_usage(rowwise_usage=False)
 
             # Weight with column-wise usage is needed for dgrad GEMM.
-            if inp.requires_grad:
-                if isinstance(weightmat, QuantizedTensor):
-                    weightmat.update_usage(columnwise_usage=True)
+            if isinstance(weightmat, QuantizedTensor):
+                weightmat.update_usage(columnwise_usage=True)
 
             if cpu_offloading:
                 if fp8 and weightmat is not None:

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -415,7 +415,7 @@ class _LayerNormMLP(torch.autograd.Function):
         )
 
         # Weight with column-wise usage is needed for dgrad GEMM.
-        if is_grad_enabled and inp.requires_grad:
+        if is_grad_enabled:
             if isinstance(fc1_weight_final, QuantizedTensor):
                 fc1_weight_final.update_usage(columnwise_usage=True)
             if isinstance(fc2_weight_final, QuantizedTensor):