create-caption

Paused

nroggendorff commited on Nov 15

Commit

dc942f1

verified ·

1 Parent(s): 87153fd

Update train.py

Files changed (1) hide show

train.py CHANGED Viewed

@@ -102,7 +102,7 @@ else:
 num_gpus = torch.cuda.device_count()
 models = [load_model(device_id=i) for i in range(num_gpus)]
-batch_size = 8
 shard_size = len(ds) // num_gpus

 num_gpus = torch.cuda.device_count()
 models = [load_model(device_id=i) for i in range(num_gpus)]
+batch_size = 32
 shard_size = len(ds) // num_gpus