minor change

pytorch · Nov 24, 2021 · 157da8b · 157da8b
1 parent db7b6db
commit 157da8b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torchtext/models/roberta/model.py b/torchtext/models/roberta/model.py
@@ -67,7 +67,7 @@ def __init__(
  return_all_layers=False,
  )
 
- self.project = None
+ self.project = nn.Identity()
  if projection_dim is not None:
  self.project = ProjectionLayer(embed_dim=embedding_dim, projection_dim=projection_dim, dropout=projection_dropout)
 
@@ -83,8 +83,7 @@ def forward(self, tokens: Tensor, mask: Optional[Tensor] = None) -> Tensor:
  if mask is not None:
  output = output[mask.to(torch.bool), :]
 
- if self.project is not None:
- output = self.project(output)
+ output = self.project(output)
 
  return output