float32 then? probably attn needs to be fp32?

hydra_node/models.py

@@ -124,7 +124,7 @@ class DalleMiniModel(nn.Module):
         self.config = config
         self.model = MinDalle(
             models_root=config.model_path,
-            dtype=torch.bfloat16,
+            dtype=torch.float32,
             device='cuda',
             is_mega=True, 
             is_reusable=True