bfloat16 because it nans

42e478be · novelailab · 1a6751df · 42e478be
Commit 42e478be authored Aug 04, 2022 by novelailab
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

hydra_node/models.py hydra_node/models.py +1 -1

No files found.
--- a/hydra_node/models.py
+++ b/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.float16,
+            dtype=torch.bfloat16,
            device='cuda',
            is_mega=True, 
            is_reusable=True