Spaces:

flax-community
/

dalle-mini

Running

boris commited on Dec 20, 2021

Commit

61c93f2

1 Parent(s): b257ca8

fix: update model name

Files changed (1) hide show

tools/train/train.py CHANGED Viewed

@@ -398,7 +398,7 @@ def main():
         artifact_dir = artifact.download()
         # load model
-        model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)
         # avoid OOM on TPU: see https://github.com/google/flax/issues/1658
         print(model.params)

         artifact_dir = artifact.download()
         # load model
+        model = DalleBart.from_pretrained(artifact_dir)
         # avoid OOM on TPU: see https://github.com/google/flax/issues/1658
         print(model.params)