fix: fix the type of the dilation argument

megvii-research · Nov 25, 2021 · a5418ea · a5418ea
1 parent e9f5235
commit a5418ea
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/main.py b/main.py
@@ -128,6 +128,14 @@ def main(args):
 
  if args.frozen_weights is not None:
  assert args.masks, "Frozen training is meant for segmentation only"
+
+ if isinstance(args.dilation,str):
+ if args.dilation.lower() == 'false':
+ args.dilation = False
+ elif args.dilation.lower() == 'true':
+ args.dilation = True
+ else:
+ raise ValueError("The dilation should be True or False")
  print(args)
 
  device = torch.device(args.device)