From 58e7cc552b33e724479ff6bb07df0cf3bb2240c1 Mon Sep 17 00:00:00 2001
From: Engin Eren <engin.eren@desy.de>
Date: Thu, 19 May 2022 15:29:53 +0200
Subject: [PATCH] Adding nworkers..
---
wganHCAL.py | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/wganHCAL.py b/wganHCAL.py
index 08b9096..822e286 100644
--- a/wganHCAL.py
+++ b/wganHCAL.py
@@ -161,6 +161,10 @@ def parse_args():
parser.add_argument('--epochs', type=int, default=1, metavar='N',
help='number of epochs to train (default: 1)')
+
+ parser.add_argument('--nworkers', type=int, default=1, metavar='N',
+ help='number of epochs to train (default: 1)')
+
parser.add_argument('--lrCrit', type=float, default=0.00001, metavar='LR',
help='learning rate Critic (default: 0.00001)')
parser.add_argument('--lrGen', type=float, default=0.0001, metavar='LR',
@@ -252,7 +256,7 @@ def run(args):
sampler = DistributedSampler(dataset, shuffle=True)
- train_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=1, drop_last=True, pin_memory=False)
+ train_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=args.nworkers, drop_last=True, pin_memory=False)
--
GitLab