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

ManukyanD · web-flow · commit de85c0917cde · 2024-01-11T20:39:33.000-08:00
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)
