Set the dtype correctly for vision GPT model (#694)

* Set the dtype correctly * Add changelog
Lightning-Universe · Jul 26, 2021 · e22b8d0 · e22b8d0
1 parent ab801aa
commit e22b8d0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -27,6 +27,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Removed momentum updating from val step and add separate val queue ([#631](https://github.com/PyTorchLightning/lightning-bolts/pull/631))
 
 
+- Fixed FP16 support with vision GPT model ([#694](https://github.com/PyTorchLightning/lightning-bolts/pull/694))
+
+
 ## [0.3.4] - 2021-06-17
 
 ### Changed

diff --git a/pl_bolts/models/vision/image_gpt/gpt2.py b/pl_bolts/models/vision/image_gpt/gpt2.py
@@ -94,7 +94,7 @@ def forward(self, x, classify=False):
         h = self.token_embeddings(x.long())
 
         # prepend sos token
-        sos = torch.ones(1, batch, self.hparams.embed_dim, device=x.device) * self.sos
+        sos = torch.ones(1, batch, self.hparams.embed_dim, device=x.device, dtype=x.dtype) * self.sos
         h = torch.cat([sos, h[:-1, :, :]], axis=0)
 
         # add positional embeddings