Bugfix in vision transformer - save class token and pos embedding (#1204

)
pytorch · Jan 12, 2024 · de85c09 · de85c09
1 parent 5a3b333
commit de85c09
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vision_transformer/main.py b/vision_transformer/main.py
@@ -49,9 +49,9 @@ def __init__(self, args):
         # Linear projection
         self.LinearProjection = nn.Linear(self.input_size, self.latent_size)
         # Class token
-        self.class_token = nn.Parameter(torch.randn(self.batch_size, 1, self.latent_size)).to(self.device)
+        self.class_token = nn.Parameter(torch.randn(self.batch_size, 1, self.latent_size).to(self.device))
         # Positional embedding
-        self.pos_embedding = nn.Parameter(torch.randn(self.batch_size, 1, self.latent_size)).to(self.device)
+        self.pos_embedding = nn.Parameter(torch.randn(self.batch_size, 1, self.latent_size).to(self.device))
 
     def forward(self, input_data):
         input_data = input_data.to(self.device)