Update build-output/torch-universal/triton_llama_mlp/mlp.py

build-output/torch-universal/triton_llama_mlp/mlp.py

@@ -152,7 +152,7 @@ class TritonLlamaMLP(nn.Module):
             down_output += self.down_proj.bias
             
         # Reshape back to original dimensions: (batch_size*seq_len, hidden_size) -> (*, hidden_size)
-        return down_output.reshape(original_shape)
+        return down_output.reshape(original_shape).to(dtype)