diff --git a/roberta.py b/roberta.py
index 346c1a7..f4d7765 100644
--- a/roberta.py
+++ b/roberta.py
@@ -52,11 +52,11 @@ class RobertaLeakyHead(nn.Module):
                 )
 
         x = self.dense_1(x)
-        x = self.leaky_relu(x)
+        x = torch.relu(x)
         x = self.dropout(x)
 
         x = self.dense_2(x)
-        x = self.leaky_relu(x)
+        x = torch.relu(x)
         x = self.dropout(x)
 
         x = self.out_proj(x)