Hi @ptrblck,
Thanks for your inputs. I moved to torch.cuda.amp.
I am using scaler = torch.cuda.amp.GradScaler().
I am doing the casting as following:
with torch.cuda.amp.autocast():
    preds = model(inputs)
    loss = criterion(preds, labels.float())
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

But I have an error in a different module for torch.einsum().
Traceback (most recent call last):
  File "main.py", line 333, in <module>
    main()
  File "main.py", line 327, in main
    train(model, train_loader, test_loader, args, criterion, optimizer, device)
  File "main.py", line 229, in train
    scaler.scale(loss).backward()
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 352, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/function.py", line 199, in apply
    return user_fn(self, *args)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/function.py", line 340, in wrapper
    outputs = fn(ctx, *args)
  File "<path>/model/pac.py", line 179, in backward
    grad_in_mul_k = torch.einsum('iomn,ojkl->ijklmn', (grad_output, weight))
  File "/opt/conda/lib/python3.8/site-packages/torch/functional.py", line 325, in einsum
    return einsum(equation, *_operands)
  File "/opt/conda/lib/python3.8/site-packages/torch/functional.py", line 327, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: expected scalar type Half but found Float

Please take a look.
Thanks!