I solved it! Woah! The error was misleading and was pointing to
abs_ = t.sqrt(re**2 + im**2)
While in reality, the error is in:
abs_[mask] = -1 * abs_[mask]

That was the inplace opeartion. I had to create a clone of that tensor and modify the clone like this:
clone = abs_.clone()
clone[mask] = -1 * abs_[mask]

Is there a more efficient way? cause abs_ is wasted in my solution. Can the clone be a detached tensor and I go like abs_ = Dtached[mask] ? IDK.