Update csrc/layernorm_kernels.cu

Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Author: ProExpertProg

csrc/layernorm_kernels.cu

@@ -380,6 +380,7 @@ void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
                         torch::Tensor& residual,  // [..., hidden_size]
                         torch::Tensor& weight,    // [hidden_size]
                         double epsilon) {
+  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
   TORCH_CHECK(input.scalar_type() == residual.scalar_type());
   TORCH_CHECK(residual.is_contiguous());
   TORCH_CHECK(weight.is_contiguous());