Skip to content

GitLab

"...text-generation-inference.git" did not exist on "b5087c4f4ece30decb739ab311448e358f6108bb"
Unverified Commit 1bd4c9e9 authored by YiYi Xu's avatar YiYi Xu Committed by GitHub
Browse files

remvoe one line as requested by gc team (#3077) 

remvoe one line
parent eb2ef316
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 4 deletions
examples/text_to_image/train_text_to_image_flax.py
View file @ 1bd4c9e9
...@@ -340,11 +340,10 @@ def main(): ...@@ -340,11 +340,10 @@ def main():
return examples return examples
if jax.process_index() == 0: if args.max_train_samples is not None:
if args.max_train_samples is not None: dataset["train"] = dataset["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
dataset["train"] = dataset["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
# Set the training transforms # Set the training transforms
train_dataset = dataset["train"].with_transform(preprocess_train) train_dataset = dataset["train"].with_transform(preprocess_train)
def collate_fn(examples): def collate_fn(examples):
pixel_values = torch.stack([example["pixel_values"] for example in examples]) pixel_values = torch.stack([example["pixel_values"] for example in examples])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment