fix learning rate arg passing

NVIDIA-AI-IOT · Sep 20, 2023 · 6536336 · 6536336
1 parent 369b0f0
commit 6536336
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nanosam/tools/train.py b/nanosam/tools/train.py
@@ -61,7 +61,7 @@
     else:
         raise RuntimeError(f"Unsupported loss function {args.loss}")
 
-    optimizer = torch.optim.Adam(image_encoder_cnn.parameters(), lr=3e-4)
+    optimizer = torch.optim.Adam(image_encoder_cnn.parameters(), lr=args.learning_rate)
 
     dataset = ImageFolder(args.images)