Train efficientnet should respect NUM environment variable (tinygrad#122

) Co-authored-by: holonomicjl <[email protected]>
mankasto · Nov 17, 2020 · 5797e63 · 5797e63
1 parent ee99d01
commit 5797e63
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/train_efficientnet.py b/examples/train_efficientnet.py
@@ -6,7 +6,7 @@
 
 if __name__ == "__main__":
   Tensor.default_gpu = os.getenv("GPU") is not None
-  model = EfficientNet()
+  model = EfficientNet(int(os.getenv("NUM", "0")))
 
   BS = 4