Make CausalBias a torch.Tensor subclass again (#121358)
# Summary
This was removed in #116071 in order to enable compile support and re-adding this seems to still work with compile
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121358
Approved by: https://github.com/Skylion007, https://github.com/cpuhrsch
