Update autoregressive/models/gpt_t2i.py

autoregressive/models/gpt_t2i.py

@@ -430,7 +430,7 @@ class Transformer(nn.Module):
                 token_embeddings = self.cls_embedding(cond_idx, train=self.training)
                 token_embeddings = token_embeddings[:,:self.cls_token_num]
                 if condition is not None:
-                    condition_embeddings = self.condition_mlp(condition.to(torch.bfloat16),train=self.training)
+                    condition_embeddings = self.condition_mlp(condition,train=self.training)#.to(torch.bfloat16),train=self.training)
                     self.condition_token = condition_embeddings
                     
             else: # decode_n_tokens(kv cache) in inference