fix mini_batch_size being referenced before assignment (#25) (#26)

openai · Jan 23, 2024 · 7055122 · 7055122
1 parent 2d5c21a
commit 7055122
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/weak_to_strong/train.py b/weak_to_strong/train.py
@@ -231,9 +231,11 @@ def maybe_load_model(model):
  ).to("cuda")
  already_trained = maybe_load_model(model)
  # data parallel: currently not supported with model parallel
+
+ minibatch_size = min(minibatch_size_per_device * torch.cuda.device_count(), batch_size)
+
  if torch.cuda.device_count() > 1:
  model = torch.nn.DataParallel(model, output_device=0)
- minibatch_size = min(minibatch_size_per_device * torch.cuda.device_count(), batch_size)
  print(
  "Using",
  torch.cuda.device_count(),