How to use PyTorch Lightning’s built-in TPU supportHow to use PyTorch Lightning’s built-in TPU supportDeveloper Advocacy Manager

After the notebook has been provisioned, select OPEN JUPYTERLAB to access the JupyterLab environment. If you’d like to access the sample for this tutorial, you can open a new terminal (File > New > Terminal), and then run:

The left sidebar will refresh after a few moments. You’ll find the sample within ai-platform samples > notebooks > samples > pytorch > lightning.

Setup your TPU node

Next, let’s provision a TPU node that we can use from our notebook instance. From the Cloud Console, go to Compute Engine > TPUs. Select Create TPU Node, and then choose a name of your choice. Then, select a Zone and TPU type, keeping in mind that TPU availability varies per region.

Make sure to select a TPU software version that matches the version you selected for your notebook, in this case pytorch-1.8. Also, for the purposes of this tutorial, select datalab-network for the Network, so that you can access your TPU directly from the notebook instance without configuring any networking settings.