Skip to content

Commit

Permalink
updates
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Jun 14, 2019
1 parent 8f60924 commit bb36820
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,11 +140,12 @@ def train(
# plt.savefig('LR.png', dpi=300)

# Dataset
rectangular_training = False
dataset = LoadImagesAndLabels(train_path,
img_size,
batch_size,
augment=True,
rect=False)
rect=rectangular_training)

# Initialize distributed training
if torch.cuda.device_count() > 1:
Expand All @@ -156,7 +157,7 @@ def train(
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=opt.num_workers,
shuffle=True, # disable rectangular training if True
shuffle=not rectangular_training, # Shuffle=True unless rectangular training is used
pin_memory=True,
collate_fn=dataset.collate_fn)

Expand Down

0 comments on commit bb36820

Please sign in to comment.