From 5797e63d9b314b0ea0b39e5c9554eae708d4acbb Mon Sep 17 00:00:00 2001
From: adamritter <58403584+adamritter@users.noreply.github.com>
Date: Tue, 17 Nov 2020 04:02:31 +0000
Subject: [PATCH] Train efficientnet should respect NUM environment variable
 (#122)

Co-authored-by: holonomicjl <58403584+holonomicjl@users.noreply.github.com>
---
 examples/train_efficientnet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/train_efficientnet.py b/examples/train_efficientnet.py
index dfb0a2b921..090e74f1ab 100644
--- a/examples/train_efficientnet.py
+++ b/examples/train_efficientnet.py
@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor
 
 if __name__ == "__main__":
   Tensor.default_gpu = os.getenv("GPU") is not None
-  model = EfficientNet()
+  model = EfficientNet(int(os.getenv("NUM", "0")))
 
   BS = 4