doc

README.md

@@ -72,7 +72,7 @@ As it turns out, 90% of what you need for neural networks are a decent autograd/
 Throw in an optimizer, a data loader, and some compute, and you have all you need.
 
 ```python
-from tinygrad import Tensor, nn
+from tinygrad import Tensor, nn, Context
 
 class LinearNet:
   def __init__(self):
@@ -86,7 +86,7 @@ optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
 
 x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader
 
-with Tensor.train():
+with Context(TRAINING=1):
   for i in range(10):
     optim.zero_grad()
     loss = model(x).sparse_categorical_crossentropy(y).backward()

docs/quickstart.md

@@ -165,13 +165,14 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.
 
-We use `with Tensor.train()` to set the internal flag `Tensor.training` to `True` during training.
+We use `with Context(TRAINING=1)` to set the internal flag `Tensor.training` to `True` during training.
 Upon exit, the flag is restored to its previous value by the context manager.
 
 ```python
+from tinygrad import Context
 X_train, Y_train, X_test, Y_test = fetch_mnist()
 
-with Tensor.train():
+with Context(TRAINING=1):
   for step in range(1000):
     # random sample a batch
     samp = np.random.randint(0, X_train.shape[0], size=(64))