[Doc] Add Nemotron to supported model docs (#6843)

2024-07-26 17:32:44 -04:00
parent 3bbb4936dc
commit 281977bd6e
2 changed files with 5 additions and 3 deletions
--- a/vllm/model_executor/layers/activation.py
+++ b/vllm/model_executor/layers/activation.py
@@ -166,9 +166,7 @@ class ReLUSquaredActivation(CustomOp):

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
-        relu_applied = nn.functional.relu(x)
-        squared = torch.square(relu_applied)
-        return squared
+        return torch.square(F.relu(x))

    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
        return self.forward_native(x)