Train efficientnet should respect NUM environment variable (#122)

Co-authored-by: holonomicjl <[email protected]>
2020-11-16 20:02:31 -08:00
parent ee99d016e9
commit 5797e63d9b
1 changed files with 1 additions and 1 deletions
@@ -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