Update evaluation/generation/generate.py

younesbelkada · web-flow · commit 719180277936 · 2022-06-28T11:12:22.000+02:00
diff --git a/evaluation/generation/generate.py b/evaluation/generation/generate.py
@@ -41,7 +41,7 @@ def main():
         device_map="auto" if args.parallelize else None,
         torch_dtype=torch.bfloat16,
         revision="gs{}".format(args.global_step) if args.global_step else None
-        offload_floder=args.offload_folder,
+        offload_folder=args.offload_folder is args.parallelize else None,
     )
     print(f"Loaded model in {datetime.datetime.now() - start}")
 

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ def main():`
`41`	`41`	`device_map="auto" if args.parallelize else None,`
`42`	`42`	`torch_dtype=torch.bfloat16,`
`43`	`43`	`revision="gs{}".format(args.global_step) if args.global_step else None`
`44`		`- offload_floder=args.offload_folder,`
	`44`	`+ offload_folder=args.offload_folder is args.parallelize else None,`
`45`	`45`	`)`
`46`	`46`	`print(f"Loaded model in {datetime.datetime.now() - start}")`
`47`	`47`