[BUG] Assertion failed: t.data.dptr != nullptr. Input x is not allocated!

[alexdremov]

While running RMS norm, I got the following exception:

/workspace/TransformerEngine/transformer_engine/common/transformer_engine.cpp:39 in function CheckInputTensor: Assertion failed: t.data.dptr != nullptr. Input x is not allocated!
  File "/usr/local/lib/python3.11/dist-packages/transformer_engine/pytorch/module/rmsnorm.py", line 50, in forward
    rmsnorm_out, rsigma = tex.rmsnorm_fwd(inputmat, rmsnorm_weight,
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/autograd/function.py", line 553, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

However, clearly, all input tensors are allocated — I verified this with debugger