Fix linear layer replacement
<!-- pylint-output-strict-mode-true -->

beep boop 🤖: 🚨 The following files must be fixed before merge!

---

Your code was analyzed with PyLint. The following annotations have been identified:

```

------------------------------------
Your code has been rated at 10.00/10
```

---

Thank you for improving NeMo's documentation! [🤖]: Hi @oyilmaz-nvidia 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/11945087930) for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g
# What does this PR do ?

Fixes the bug with TE linear layer replacement