<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 <a href="https
Step1:  Setup lightning
Step2:  Iris
Step3:  CIFAR
<ASSISTANT_TASK:>
Python Code:
import matplotlib.pyplot as plt
import numpy as np
import os
assert os.environ["COLAB_TPU_ADDR"], "Make sure to select TPU from Edit > Notebook settings > Hardware accelerator"
#!pip install -q cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8-cp37-cp37m-linux_x86_64.whl
!pip install -q cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl
Explanation: <a href="https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/logreg_tpu_pytorch_lightning_bolts.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Logistic regression on MNIST using TPUs and PyTorch Lightning
Code is from 
https://lightning-bolts.readthedocs.io/en/latest/introduction_guide.html#logistic-regression
Setup TPU
Be sure to select Runtime=TPU in the drop-down menu!
See
https://colab.sandbox.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb#scrollTo=3P6b3uqfzpDI
See also 
https://pytorch-lightning.readthedocs.io/en/latest/notebooks/lightning_examples/cifar10-baseline.html#
End of explanation
#!pip install -q lightning-bolts
!pip install --quiet torchmetrics lightning-bolts torchvision torch pytorch-lightning
from pl_bolts.models.regression import LogisticRegression
import pytorch_lightning as pl
from pl_bolts.datamodules import MNISTDataModule, FashionMNISTDataModule, CIFAR10DataModule, ImagenetDataModule
Explanation: Setup lightning
End of explanation
from sklearn.datasets import load_iris
from pl_bolts.datamodules import SklearnDataModule
import pytorch_lightning as pl
# use any numpy or sklearn dataset
X, y = load_iris(return_X_y=True)
dm = SklearnDataModule(X, y, batch_size=12)
# build model
model = LogisticRegression(input_dim=4, num_classes=3)
# fit
trainer = pl.Trainer(tpu_cores=8)
trainer.fit(model, train_dataloader=dm.train_dataloader(), val_dataloaders=dm.val_dataloader())
trainer.test(test_dataloaders=dm.test_dataloader())
Explanation: Iris
End of explanation
# create dataset
# dm = MNISTDataModule(num_workers=0, data_dir='data')
dm = CIFAR10DataModule(num_workers=0, data_dir="data")
dm.prepare_data()  # force download now
print(dm.size())
print(dm.num_classes)
ndims = np.prod(dm.size())
nclasses = dm.num_classes
print([ndims, nclasses, ndims * nclasses])
model = LogisticRegression(input_dim=ndims, num_classes=nclasses, learning_rate=0.001)
print(model)
trainer = pl.Trainer(tpu_cores=8, max_epochs=2)
# trainer = pl.Trainer(max_epochs=2)
trainer.fit(model, datamodule=dm)
trainer.test(model, test_dataloaders=dm.val_dataloader())
Explanation: CIFAR
End of explanation
<END_TASK>