```
VERSION = "20200707" #"nightly"  #"20200515" @param ["1.5" , "20200325", "nightly"]
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py
```

      % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                     Dload  Upload   Total   Spent    Left  Speed
    100  5115  100  5115    0     0  53842      0 --:--:-- --:--:-- --:--:-- 53842
    Updating... This may take around 2 minutes.
    Updating TPU runtime to pytorch-dev20200515 ...
    Collecting cloud-tpu-client
      Downloading https://files.pythonhosted.org/packages/56/9f/7b1958c2886db06feb5de5b2c191096f9e619914b6c31fdf93999fdbbd8b/cloud_tpu_client-0.10-py3-none-any.whl
    Collecting google-api-python-client==1.8.0
      Downloading https://files.pythonhosted.org/packages/9a/b4/a955f393b838bc47cbb6ae4643b9d0f90333d3b4db4dc1e819f36aad18cc/google_api_python_client-1.8.0-py3-none-any.whl (57kB)
         |████████████████████████████████| 61kB 2.6MB/s
    Requirement already satisfied: oauth2client in /usr/local/lib/python3.6/dist-packages (from cloud-tpu-client) (4.1.3)
    Requirement already satisfied: google-api-core<2dev,>=1.13.0 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (1.16.0)
    Requirement already satisfied: httplib2<1dev,>=0.9.2 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (0.17.4)
    Requirement already satisfied: google-auth>=1.4.1 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (1.17.2)
    Requirement already satisfied: six<2dev,>=1.6.1 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (1.12.0)
    Requirement already satisfied: google-auth-httplib2>=0.0.3 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (0.0.4)
    Requirement already satisfied: uritemplate<4dev,>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from google-api-python-client==1.8.0->cloud-tpu-client) (3.0.1)
    Requirement already satisfied: pyasn1>=0.1.7 in /usr/local/lib/python3.6/dist-packages (from oauth2client->cloud-tpu-client) (0.4.8)
    Requirement already satisfied: rsa>=3.1.4 in /usr/local/lib/python3.6/dist-packages (from oauth2client->cloud-tpu-client) (4.6)
    Requirement already satisfied: pyasn1-modules>=0.0.5 in /usr/local/lib/python3.6/dist-packages (from oauth2client->cloud-tpu-client) (0.2.8)
    Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (1.52.0)
    Requirement already satisfied: pytz in /usr/local/lib/python3.6/dist-packages (from google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (2018.9)
    Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /usr/local/lib/python3.6/dist-packages (from google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (2.23.0)
    Requirement already satisfied: protobuf>=3.4.0 in /usr/local/lib/python3.6/dist-packages (from google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (3.12.2)
    Requirement already satisfied: setuptools>=34.0.0 in /usr/local/lib/python3.6/dist-packages (from google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (49.1.0)
    Uninstalling torch-1.5.1+cu101:
    Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth>=1.4.1->google-api-python-client==1.8.0->cloud-tpu-client) (4.1.1)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (1.24.3)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (3.0.4)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (2.10)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.13.0->google-api-python-client==1.8.0->cloud-tpu-client) (2020.6.20)
    Installing collected packages: google-api-python-client, cloud-tpu-client
      Found existing installation: google-api-python-client 1.7.12
        Uninstalling google-api-python-client-1.7.12:
          Successfully uninstalled google-api-python-client-1.7.12
    Successfully installed cloud-tpu-client-0.10 google-api-python-client-1.8.0
    Done updating TPU runtime
      Successfully uninstalled torch-1.5.1+cu101
    Uninstalling torchvision-0.6.1+cu101:
      Successfully uninstalled torchvision-0.6.1+cu101
    Copying gs://tpu-pytorch/wheels/torch-nightly+20200515-cp36-cp36m-linux_x86_64.whl...
    \ [1 files][ 91.0 MiB/ 91.0 MiB]
    Operation completed over 1 objects/91.0 MiB.
    Copying gs://tpu-pytorch/wheels/torch_xla-nightly+20200515-cp36-cp36m-linux_x86_64.whl...
    \ [1 files][119.5 MiB/119.5 MiB]
    Operation completed over 1 objects/119.5 MiB.
    Copying gs://tpu-pytorch/wheels/torchvision-nightly+20200515-cp36-cp36m-linux_x86_64.whl...
    / [1 files][  2.3 MiB/  2.3 MiB]
    Operation completed over 1 objects/2.3 MiB.
    Processing ./torch-nightly+20200515-cp36-cp36m-linux_x86_64.whl
    Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch==nightly+20200515) (0.16.0)
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==nightly+20200515) (1.18.5)
    ERROR: fastai 1.0.61 requires torchvision, which is not installed.
    Installing collected packages: torch
    Successfully installed torch-1.6.0a0+bf2bbd9
    Processing ./torch_xla-nightly+20200515-cp36-cp36m-linux_x86_64.whl
    Installing collected packages: torch-xla
    Successfully installed torch-xla-1.6+2b2085a
    Processing ./torchvision-nightly+20200515-cp36-cp36m-linux_x86_64.whl
    Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200515) (7.0.0)
    Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200515) (1.6.0a0+bf2bbd9)
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200515) (1.18.5)
    Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch->torchvision==nightly+20200515) (0.16.0)
    Installing collected packages: torchvision
    Successfully installed torchvision-0.7.0a0+a6073f0
    Reading package lists... Done
    Building dependency tree
    Reading state information... Done
    The following package was automatically installed and is no longer required:
      libnvidia-common-440
    Use 'apt autoremove' to remove it.
    The following NEW packages will be installed:
      libomp5
    0 upgraded, 1 newly installed, 0 to remove and 35 not upgraded.
    Need to get 234 kB of archives.
    After this operation, 774 kB of additional disk space will be used.
    Get:1 http://archive.ubuntu.com/ubuntu bionic/universe amd64 libomp5 amd64 5.0.1-1 [234 kB]
    Fetched 234 kB in 1s (374 kB/s)
    Selecting previously unselected package libomp5:amd64.
    (Reading database ... 144465 files and directories currently installed.)
    Preparing to unpack .../libomp5_5.0.1-1_amd64.deb ...
    Unpacking libomp5:amd64 (5.0.1-1) ...
    Setting up libomp5:amd64 (5.0.1-1) ...
    Processing triggers for libc-bin (2.27-3ubuntu1) ...
    /sbin/ldconfig.real: /usr/local/lib/python3.6/dist-packages/ideep4py/lib/libmkldnn.so.0 is not a symbolic link

```
!pip install https://github.com/fastai/fastcore/archive/master.zip
!pip install https://github.com/fastai/fastai2/archive/master.zip
```

    Collecting https://github.com/fastai/fastcore/archive/master.zip
      Downloading https://github.com/fastai/fastcore/archive/master.zip
         - 3.5MB 2.1MB/s
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from fastcore==0.1.18) (1.18.5)
    Requirement already satisfied: dataclasses>='0.7' in /usr/local/lib/python3.6/dist-packages (from fastcore==0.1.18) (0.7)
    Building wheels for collected packages: fastcore
      Building wheel for fastcore (setup.py) ... done
      Created wheel for fastcore: filename=fastcore-0.1.18-cp36-none-any.whl size=28891 sha256=36950ebbe97c27ebcc38471ed8b876294aa4547427aa1e604b3cabdaf90ed9eb
      Stored in directory: /tmp/pip-ephem-wheel-cache-1e3_h2zu/wheels/68/af/c5/c8a9f7370515ab9b237b3fd59c04a0e7c5d4dedc7a6768a772
    Successfully built fastcore
    Installing collected packages: fastcore
    Successfully installed fastcore-0.1.18
    Collecting https://github.com/fastai/fastai2/archive/master.zip
      Downloading https://github.com/fastai/fastai2/archive/master.zip
         - 121.7MB 795kB/s
    Requirement already satisfied: fastcore in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (0.1.18)
    Requirement already satisfied: torch>=1.3.0 in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (1.6.0a0+bf2bbd9)
    Requirement already satisfied: torchvision>=0.5 in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (0.7.0a0+a6073f0)
    Requirement already satisfied: matplotlib in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (3.2.2)
    Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (1.0.5)
    Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (2.23.0)
    Requirement already satisfied: pyyaml in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (3.13)
    Requirement already satisfied: fastprogress>=0.1.22 in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (0.2.3)
    Requirement already satisfied: pillow in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (7.0.0)
    Requirement already satisfied: scikit-learn in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (0.22.2.post1)
    Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (1.4.1)
    Requirement already satisfied: spacy in /usr/local/lib/python3.6/dist-packages (from fastai2==0.0.18) (2.2.4)
    Requirement already satisfied: dataclasses>='0.7'; python_version < "3.7" in /usr/local/lib/python3.6/dist-packages (from fastcore->fastai2==0.0.18) (0.7)
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from fastcore->fastai2==0.0.18) (1.18.5)
    Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch>=1.3.0->fastai2==0.0.18) (0.16.0)
    Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib->fastai2==0.0.18) (0.10.0)
    Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->fastai2==0.0.18) (1.2.0)
    Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->fastai2==0.0.18) (2.4.7)
    Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->fastai2==0.0.18) (2.8.1)
    Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas->fastai2==0.0.18) (2018.9)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->fastai2==0.0.18) (3.0.4)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->fastai2==0.0.18) (1.24.3)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->fastai2==0.0.18) (2.10)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->fastai2==0.0.18) (2020.6.20)
    Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->fastai2==0.0.18) (0.16.0)
    Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (1.0.2)
    Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (0.7.0)
    Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (0.4.1)
    Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (4.41.1)
    Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (3.0.2)
    Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (2.0.3)
    Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (1.0.0)
    Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (7.4.0)
    Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (49.1.0)
    Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (1.0.2)
    Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy->fastai2==0.0.18) (1.1.3)
    Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from cycler>=0.10->matplotlib->fastai2==0.0.18) (1.12.0)
    Requirement already satisfied: importlib-metadata>=0.20; python_version < "3.8" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy->fastai2==0.0.18) (1.7.0)
    Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < "3.8"->catalogue<1.1.0,>=0.0.7->spacy->fastai2==0.0.18) (3.1.0)
    Building wheels for collected packages: fastai2
      Building wheel for fastai2 (setup.py) ... done
      Created wheel for fastai2: filename=fastai2-0.0.18-cp36-none-any.whl size=175594 sha256=2d921efd963d6f8fb2b1d7a710bf76a9be354aea3ecd090db9583f82a05ed1b3
      Stored in directory: /tmp/pip-ephem-wheel-cache-1od2hht4/wheels/b2/28/3b/7a38644c6129e851e0db09b594949028266a62f4cb234fd1f3
    Successfully built fastai2
    Installing collected packages: fastai2
    Successfully installed fastai2-0.0.18

```
!pip install py-heat-magic
```

    Collecting py-heat-magic
      Downloading https://files.pythonhosted.org/packages/72/15/c5c56c91077ecb42a6415364cff41acb08a4f696f0261b3fc4025856d85c/py-heat-magic-0.0.2.tar.gz
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (1.18.5)
    Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (1.4.1)
    Requirement already satisfied: matplotlib in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (3.2.2)
    Requirement already satisfied: ipython in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (5.5.0)
    Requirement already satisfied: jupyter in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (1.0.0)
    Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (1.0.5)
    Requirement already satisfied: sympy in /usr/local/lib/python3.6/dist-packages (from py-heat-magic) (1.1.1)
    Collecting nose
      Downloading https://files.pythonhosted.org/packages/15/d8/dd071918c040f50fa1cf80da16423af51ff8ce4a0f2399b7bf8de45ac3d9/nose-1.3.7-py3-none-any.whl (154kB)
         |████████████████████████████████| 163kB 3.5MB/s
    Collecting py-heat
      Downloading https://files.pythonhosted.org/packages/6b/e3/776c6f4f18eafbc39b44ad47160414077bfd64183cfd2c5b61ea43dd12b6/py-heat-0.0.6.tar.gz
    Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->py-heat-magic) (1.2.0)
    Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->py-heat-magic) (2.4.7)
    Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->py-heat-magic) (2.8.1)
    Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib->py-heat-magic) (0.10.0)
    Requirement already satisfied: simplegeneric>0.8 in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (0.8.1)
    Requirement already satisfied: decorator in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (4.4.2)
    Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (4.3.3)
    Requirement already satisfied: pickleshare in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (0.7.5)
    Requirement already satisfied: prompt-toolkit<2.0.0,>=1.0.4 in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (1.0.18)
    Requirement already satisfied: pexpect; sys_platform != "win32" in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (4.8.0)
    Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (49.1.0)
    Requirement already satisfied: pygments in /usr/local/lib/python3.6/dist-packages (from ipython->py-heat-magic) (2.1.3)
    Requirement already satisfied: nbconvert in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (5.6.1)
    Requirement already satisfied: notebook in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (5.2.2)
    Requirement already satisfied: jupyter-console in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (5.2.0)
    Requirement already satisfied: qtconsole in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (4.7.5)
    Requirement already satisfied: ipykernel in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (4.10.1)
    Requirement already satisfied: ipywidgets in /usr/local/lib/python3.6/dist-packages (from jupyter->py-heat-magic) (7.5.1)
    Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas->py-heat-magic) (2018.9)
    Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.6/dist-packages (from sympy->py-heat-magic) (1.1.0)
    Collecting pprofile
      Downloading https://files.pythonhosted.org/packages/dd/db/198fd60559f77c5334f520e2a94cde5c8c6f5f4e65f82f5b22269e6e5ff2/pprofile-2.0.5.tar.gz (54kB)
         |████████████████████████████████| 61kB 5.5MB/s
    Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.6/dist-packages (from python-dateutil>=2.1->matplotlib->py-heat-magic) (1.12.0)
    Requirement already satisfied: ipython-genutils in /usr/local/lib/python3.6/dist-packages (from traitlets>=4.2->ipython->py-heat-magic) (0.2.0)
    Requirement already satisfied: wcwidth in /usr/local/lib/python3.6/dist-packages (from prompt-toolkit<2.0.0,>=1.0.4->ipython->py-heat-magic) (0.2.5)
    Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.6/dist-packages (from pexpect; sys_platform != "win32"->ipython->py-heat-magic) (0.6.0)
    Requirement already satisfied: mistune<2,>=0.8.1 in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (0.8.4)
    Requirement already satisfied: nbformat>=4.4 in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (5.0.7)
    Requirement already satisfied: testpath in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (0.4.4)
    Requirement already satisfied: defusedxml in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (0.6.0)
    Requirement already satisfied: jinja2>=2.4 in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (2.11.2)
    Requirement already satisfied: entrypoints>=0.2.2 in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (0.3)
    Requirement already satisfied: bleach in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (3.1.5)
    Requirement already satisfied: jupyter-core in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (4.6.3)
    Requirement already satisfied: pandocfilters>=1.4.1 in /usr/local/lib/python3.6/dist-packages (from nbconvert->jupyter->py-heat-magic) (1.4.2)
    Requirement already satisfied: terminado>=0.3.3; sys_platform != "win32" in /usr/local/lib/python3.6/dist-packages (from notebook->jupyter->py-heat-magic) (0.8.3)
    Requirement already satisfied: jupyter-client in /usr/local/lib/python3.6/dist-packages (from notebook->jupyter->py-heat-magic) (5.3.5)
    Requirement already satisfied: tornado>=4 in /usr/local/lib/python3.6/dist-packages (from notebook->jupyter->py-heat-magic) (4.5.3)
    Requirement already satisfied: qtpy in /usr/local/lib/python3.6/dist-packages (from qtconsole->jupyter->py-heat-magic) (1.9.0)
    Requirement already satisfied: pyzmq>=17.1 in /usr/local/lib/python3.6/dist-packages (from qtconsole->jupyter->py-heat-magic) (19.0.1)
    Requirement already satisfied: widgetsnbextension~=3.5.0 in /usr/local/lib/python3.6/dist-packages (from ipywidgets->jupyter->py-heat-magic) (3.5.1)
    Requirement already satisfied: jsonschema!=2.5.0,>=2.4 in /usr/local/lib/python3.6/dist-packages (from nbformat>=4.4->nbconvert->jupyter->py-heat-magic) (2.6.0)
    Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.6/dist-packages (from jinja2>=2.4->nbconvert->jupyter->py-heat-magic) (1.1.1)
    Requirement already satisfied: webencodings in /usr/local/lib/python3.6/dist-packages (from bleach->nbconvert->jupyter->py-heat-magic) (0.5.1)
    Requirement already satisfied: packaging in /usr/local/lib/python3.6/dist-packages (from bleach->nbconvert->jupyter->py-heat-magic) (20.4)
    Building wheels for collected packages: py-heat-magic, py-heat, pprofile
      Building wheel for py-heat-magic (setup.py) ... done
      Created wheel for py-heat-magic: filename=py_heat_magic-0.0.2-cp36-none-any.whl size=3161 sha256=f98d9a7fdf2188985e1a1d9825d997f3ba748b2359e86bbf2e48674407327521
      Stored in directory: /root/.cache/pip/wheels/a0/4c/14/8858096734820b8a1eece9f174fe4a96d199ba247a671dc7ed
      Building wheel for py-heat (setup.py) ... done
      Created wheel for py-heat: filename=py_heat-0.0.6-py2.py3-none-any.whl size=7801 sha256=74513465d758625bd4d9c4348b5d9604068b6296111eaa919ebed7cf5c2ebf18
      Stored in directory: /root/.cache/pip/wheels/9d/3d/b7/e191bd2f2d8c0a6fc5008a87cdf69d6bc009ee3e22149c3ef3
      Building wheel for pprofile (setup.py) ... done
      Created wheel for pprofile: filename=pprofile-2.0.5-cp36-none-any.whl size=34838 sha256=ba4dd228bb5da0055d2151f28686e85c73f11a3e69248e5bd220e23e790eda39
      Stored in directory: /root/.cache/pip/wheels/b7/7e/dc/fa6663a43280716fe9a4fedfcfaae1e22fc26c701afff093c6
    Successfully built py-heat-magic py-heat pprofile
    Installing collected packages: nose, pprofile, py-heat, py-heat-magic
    Successfully installed nose-1.3.7 pprofile-2.0.5 py-heat-0.0.6 py-heat-magic-0.0.2

```
%load_ext heat
```

```
%%heat
#import pdb
print("---------------------------------------------------------- START")
#import torch
import torch_xla.core.xla_model as xm

tpu_device = xm.xla_device()
#tpu_device = torch.device('cuda:0') #torch.cuda.device("cuda:0")#.current_device()
print(f"device is {tpu_device}")
print("---------------------------------------------------------- START:1")

from fastai2.vision.all import *
path = untar_data(URLs.MNIST)
datablock = DataBlock(
    blocks=(ImageBlock,CategoryBlock),
    get_items=get_image_files,
    get_y=parent_label,
    splitter=RandomSubsetSplitter(train_sz=0.001,valid_sz=0.001),#GrandparentSplitter(train_name='training',valid_name='testing'),
#    item_tfms=Resize(28),
    #batch_tfms=aug_transforms(size=224)
    batch_tfms=aug_transforms(size=224,min_scale=0.75)
)
print("---------------------------------------------------------- START:2")

#datablock.summary(path)

# %%
printed_device=0
printed_device2=0
NUM_PIXELS=3 # 1 single channedl 3 rgb
NUM_OUTPUTS=10 #2 for MNIST_TINY 10 MNIST
class Lenet2(nn.Module):
    """Lenet with layers"""
    def __init__(self):
        super(Lenet2, self).__init__()
        self.conv1 = nn.Conv2d(NUM_PIXELS, 6, 3) # set 3 for first item if RGB
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(46656, 120) #(400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, NUM_OUTPUTS) # Only 2 outputs (3 and 7) instead of 10
    def forward(self, x):
        global printed_device
        if x.device != tpu_device:
          print(f"************** ######################### EL DEVICE ES {tpu_device} no es {x.device} ######################### **************") if printed_device < 10 else noop
          printed_device += 1
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

print("---------------------------------------------------------- START:3")

# %%
dls_tpu = datablock.dataloaders(path,bs=1, device=tpu_device)
dls_tpu.vocab, dls_tpu.show_batch()

print("---------------------------------------------------------- START:4")

# %%
def print_local(msg):
  if False: return
  print(msg)

print("---------------------------------------------------------- START:5")

class XLAOptimProxy:
    def __init__(self,opt:Optimizer):
        # TRACE: print("XLAOptimProxy#inicializando __init__")
        self.opt = opt

    def xla_step(self):
        print("------------- xla optimizer!!!!!!!! BARRIER TRYE")
        xm.optimizer_step(self.opt,barrier=True) # sync on gradient update

    def __getattr__(self,name):
        if name == '___step___': # override proxying for step method
                print_local("calling xla_step")
                return getattr(self,'xla_step')
        # proxy everything else
        #print_local(f"calling XLAOptimProxy#{name}")
        return getattr(self.opt,name)

print("---------------------------------------------------------- START:6")

@patch_to(ParamScheduler)
def _update_val(self, pct):
#        for n,f in self.scheds.items(): self.opt.set_hyper(n, f(pct))
        for n,f in self.scheds.items():
            v = f(pct)
            #print_local(f"---------------------- A f(pct) = {v}")
            self.opt.set_hyper(n, v)

print("---------------------------------------------------------- START:7")

@patch_to(ParamScheduler)
def after_batch(self):
#        for p in self.scheds.keys(): self.hps[p].append(self.opt.hypers[-1][p])
        for p in self.scheds.keys():
            v = self.opt.hypers[-1][p]
            #print_local(f"---------------------- B after_batch ParamScheduler {v}")
            self.hps[p].append(v)

print("---------------------------------------------------------- START:8 enbd STARET")

# %%
####################################################################################333

print("---------------------- STARTED class Learner")

@log_args(but='dls,model,opt_func,cbs')
class Learner():
    def __init__(self, dls, model, loss_func=None, opt_func=Adam, lr=defaults.lr, splitter=trainable_params, cbs=None,
                 metrics=None, path=None, model_dir='models', wd=None, wd_bn_bias=False, train_bn=True,
                 moms=(0.95,0.85,0.95)):
        store_attr(self, "dls,model,opt_func,lr,splitter,model_dir,wd,wd_bn_bias,train_bn,metrics,moms")
        self.training,self.create_mbar,self.logger,self.opt,self.cbs = False,True,print,None,L()
        if loss_func is None:
            loss_func = getattr(dls.train_ds, 'loss_func', None)
            assert loss_func is not None, "Could not infer loss function from the data, please pass a loss function."
        self.loss_func = loss_func
        self.path = Path(path) if path is not None else getattr(dls, 'path', Path('.'))
        self.add_cbs([(cb() if isinstance(cb, type) else cb) for cb in L(defaults.callbacks)+L(cbs)])
        self.model.to(self.dls.device)
        if hasattr(self.model, 'reset'): self.model.reset()
        self.epoch,self.n_epoch,self.loss = 0,1,tensor(0.)

    @property
    def metrics(self): return self._metrics
    @metrics.setter
    def metrics(self,v): self._metrics = L(v).map(mk_metric)

    def _grab_cbs(self, cb_cls): return L(cb for cb in self.cbs if isinstance(cb, cb_cls))
    def add_cbs(self, cbs): L(cbs).map(self.add_cb)
    def remove_cbs(self, cbs): L(cbs).map(self.remove_cb)
    def add_cb(self, cb):
        old = getattr(self, cb.name, None)
        assert not old or isinstance(old, type(cb)), f"self.{cb.name} already registered"
        cb.learn = self
        setattr(self, cb.name, cb)
        self.cbs.append(cb)
        return self

    def remove_cb(self, cb):
        if isinstance(cb, type): self.remove_cbs(self._grab_cbs(cb))
        else:
            cb.learn = None
            if hasattr(self, cb.name): delattr(self, cb.name)
            if cb in self.cbs: self.cbs.remove(cb)

    @contextmanager
    def added_cbs(self, cbs):
        self.add_cbs(cbs)
        try: yield
        finally: self.remove_cbs(cbs)

    @contextmanager
    def removed_cbs(self, cbs):
        self.remove_cbs(cbs)
        try: yield self
        finally: self.add_cbs(cbs)

    def ordered_cbs(self, event): return [cb for cb in sort_by_run(self.cbs) if hasattr(cb, event)]

    def __call__(self, event_name): L(event_name).map(self._call_one)
    def _call_one(self, event_name):
        assert hasattr(event, event_name)
        myl = []
        if False:
            # NOTE: no va
            return
        # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   SORT BY NAME")
        lll = sort_by_run(self.cbs)
        # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   SORT BY NAME {lll}")
        for cb in lll:
            # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   calling {cb}({event_name})")
            r=cb(event_name)
            # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   calling {cb}({event_name})")
            myl.append(r)
        # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   SORT BY NAME END")
        # TRACE: print("-------------------------------------------------------------------------------------------------")
        # TRACE: print("----------------------------------------------------------")
        # TRACE: print("------------------------")
        #[cb(event_name) for cb in sort_by_run(self.cbs)]

    def _bn_bias_state(self, with_bias): return bn_bias_params(self.model, with_bias).map(self.opt.state)
    def create_opt(self):
        self.opt = self.opt_func(self.splitter(self.model), lr=self.lr)
        if not self.wd_bn_bias:
            for p in self._bn_bias_state(True ): p['do_wd'] = False
        if self.train_bn:
            for p in self._bn_bias_state(False): p['force_train'] = True

    def _split(self, b):
        i = getattr(self.dls, 'n_inp', 1 if len(b)==1 else len(b)-1)
        self.xb,self.yb = b[:i],b[i:]

    def all_batches(self):
        self.n_iter = len(self.dl)
        # TRACE: print(f"Learner#all_batches@n_iter = {self.n_iter}")
        for o in enumerate(self.dl): self.one_batch(*o)

    def one_batch(self, i, b):
        # TRACE: print(f"Learner#one_batch {i} b is the batch") if i % 10000 == 0 else noop
        self.iter = i
        try:
            self._split(b);                                  self('begin_batch')
            self.pred = self.model(*self.xb);                self('after_pred')
            if len(self.yb) == 0: return
            self.loss = self.loss_func(self.pred, *self.yb); self('after_loss')
            if not self.training: return
            self.loss.backward();                            self('after_backward')
            #self.opt.step();                                 self('after_step')
            xm.optimizer_step(self.opt,barrier=True);                                 self('after_step')
            self.opt.zero_grad()
        except CancelBatchException:                         self('after_cancel_batch')
        finally:                                             self('after_batch')

    def _do_begin_fit(self, n_epoch):
        # TRACE: print("Learner#_do_begin_fit")
        self.n_epoch,self.loss = n_epoch,tensor(0.);         self('begin_fit')

    def _do_epoch_train(self):
        try:
            # TRACE: print("Learner#_do_epoch_train internal try")
            self.dl = self.dls.train;                        self('begin_train')
            self.all_batches()
        except CancelTrainException:                         self('after_cancel_train')
        finally:                                             self('after_train')

    def _do_epoch_validate(self, ds_idx=1, dl=None):
        if dl is None: dl = self.dls[ds_idx]
        try:
            # TRACE: print("Learner#_do_epoch_validate@internal try")
            self.dl = dl;                                    self('begin_validate')
            with torch.no_grad(): self.all_batches()
        except CancelValidException:                         self('after_cancel_validate')
        finally:                                             self('after_validate')

    def _end_cleanup(self):
        self.dl,self.xb,self.yb,self.pred,self.loss = None,(None,),(None,),None,None

    @log_args(but='cbs')
    def fit(self, n_epoch, lr=None, wd=None, cbs=None, reset_opt=False):
        # TRACE: print("fit starts")
        with self.added_cbs(cbs):
            # TRACE: print("fit starts with added callbacks")
            if reset_opt or not self.opt: self.create_opt()
            if wd is None: wd = self.wd
            if wd is not None: self.opt.set_hypers(wd=wd)
            # TRACE: print("Learner#fit@set hypers")
            self.opt.set_hypers(lr=self.lr if lr is None else lr)
            try:
                # TRACE: print("Learner#fit@do begin fit")
                self._do_begin_fit(n_epoch)
                # TRACE: print("Learner#fit@FOR each epoch LOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOP")
                for epoch in range(n_epoch):
                    try:
                        # TRACE: print(f"Learner#fit@epoch... begin... {epoch}")
                        self.epoch=epoch;          self('begin_epoch')
                        # TRACE: print(f"Learner#fit@epoch... train... {epoch}")
                        self._do_epoch_train()
                        # TRACE: print(f"Learner#fit@epoc validate {epoch}")
                        self._do_epoch_validate()
                    except CancelEpochException:   self('after_cancel_epoch')
                    finally:                       self('after_epoch')

            except CancelFitException:             self('after_cancel_fit')
            finally:
                # TRACE: self('after_fit')
                self._end_cleanup()
                # TRACE: print("Learner#fit@end cleanup")

    def validate(self, ds_idx=1, dl=None, cbs=None):
        if dl is None: dl = self.dls[ds_idx]
        with self.added_cbs(cbs), self.no_logging(), self.no_mbar():
            self(_before_epoch)
            self._do_epoch_validate(ds_idx, dl)
            self(_after_epoch)
        return getattr(self, 'final_record', None)

    @delegates(GatherPredsCallback.__init__)
    def get_preds(self, ds_idx=1, dl=None, with_input=False, with_decoded=False, with_loss=False, act=None,
                  inner=False, reorder=True, **kwargs):
        if dl is None: dl = self.dls[ds_idx].new(shuffled=False, drop_last=False)
        if reorder and hasattr(dl, 'get_idxs'):
            idxs = dl.get_idxs()
            dl = dl.new(get_idxs = _ConstantFunc(idxs))
        cb = GatherPredsCallback(with_input=with_input, with_loss=with_loss, **kwargs)
        ctx_mgrs = [self.no_logging(), self.added_cbs(cb), self.no_mbar()]
        if with_loss: ctx_mgrs.append(self.loss_not_reduced())
        with ExitStack() as stack:
            for mgr in ctx_mgrs: stack.enter_context(mgr)
            self(event.begin_epoch if inner else _before_epoch)
            self._do_epoch_validate(dl=dl)
            self(event.after_epoch if inner else _after_epoch)
            if act is None: act = getattr(self.loss_func, 'activation', noop)
            res = cb.all_tensors()
            pred_i = 1 if with_input else 0
            if res[pred_i] is not None:
                res[pred_i] = act(res[pred_i])
                if with_decoded: res.insert(pred_i+2, getattr(self.loss_func, 'decodes', noop)(res[pred_i]))
            if reorder and hasattr(dl, 'get_idxs'): res = nested_reorder(res, tensor(idxs).argsort())
            return tuple(res)
        self._end_cleanup()

    def predict(self, item, rm_type_tfms=None, with_input=False):
        dl = self.dls.test_dl([item], rm_type_tfms=rm_type_tfms, num_workers=0)
        inp,preds,_,dec_preds = self.get_preds(dl=dl, with_input=True, with_decoded=True)
        i = getattr(self.dls, 'n_inp', -1)
        inp = (inp,) if i==1 else tuplify(inp)
        dec = self.dls.decode_batch(inp + tuplify(dec_preds))[0]
        dec_inp,dec_targ = map(detuplify, [dec[:i],dec[i:]])
        res = dec_targ,dec_preds[0],preds[0]
        if with_input: res = (dec_inp,) + res
        return res

    def show_results(self, ds_idx=1, dl=None, max_n=9, shuffle=True, **kwargs):
        if dl is None: dl = self.dls[ds_idx].new(shuffle=shuffle)
        b = dl.one_batch()
        _,_,preds = self.get_preds(dl=[b], with_decoded=True)
        self.dls.show_results(b, preds, max_n=max_n, **kwargs)

    def show_training_loop(self):
        indent = 0
        for s in _loop:
            if s.startswith('Start'): print(f'{" "*indent}{s}'); indent += 2
            elif s.startswith('End'): indent -= 2; print(f'{" "*indent}{s}')
            else: print(f'{" "*indent} - {s:15}:', self.ordered_cbs(s))

    @contextmanager
    def no_logging(self): return replacing_yield(self, 'logger', noop)
    @contextmanager
    def no_mbar(self):    return replacing_yield(self, 'create_mbar', False)

    @contextmanager
    def loss_not_reduced(self):
        if hasattr(self.loss_func, 'reduction'): return replacing_yield(self.loss_func, 'reduction', 'none')
        else: return replacing_yield(self, 'loss_func', partial(self.loss_func, reduction='none'))

    @delegates(save_model)
    def save(self, file, **kwargs):
        file = join_path_file(file, self.path/self.model_dir, ext='.pth')
        save_model(file, self.model, getattr(self,'opt',None), **kwargs)

    @delegates(load_model)
    def load(self, file, with_opt=None, device=None, **kwargs):
        if device is None: device = self.dls.device
        if self.opt is None: self.create_opt()
        file = join_path_file(file, self.path/self.model_dir, ext='.pth')
        load_model(file, self.model, self.opt, device=device, **kwargs)
        return self

Learner.x,Learner.y = add_props(lambda i,x: detuplify((x.xb,x.yb)[i]))
###################################################################33333

def _is_instance(f, gs):
    tst = [g if type(g) in [type, 'function'] else g.__class__ for g in gs]
    for g in tst:
        if isinstance(f, g) or f==g: return True
    return False

def _is_first(f, gs):
    for o in L(getattr(f, 'run_after', None)):
        if _is_instance(o, gs): return False
    for g in gs:
        if _is_instance(f, L(getattr(g, 'run_before', None))): return False
    return True

def sort_by_run(fs):
    end = L(fs).attrgot('toward_end')
    inp,res = L(fs)[~end] + L(fs)[end], L()
    while len(inp):
        for i,o in enumerate(inp):
            if _is_first(o, inp):
                res.append(inp.pop(i))
                break
        else: raise Exception("Impossible to sort")
    # TRACE: print(f"will invoke toward_end {res}")
    return res

print("---------------------- PASSED LEARNER")

##################################################################33 torch_base.py

print("---------------------- STARTED torch_base.py")

if torch.cuda.is_available():
    if torch.cuda.current_device()==0:
        def_gpu = int(os.environ.get('DEFAULT_GPU') or 0)
        if torch.cuda.device_count()>=def_gpu: torch.cuda.set_device(def_gpu)
    torch.backends.cudnn.benchmark = True

# Cell
@delegates(plt.subplots, keep=True)
def subplots(nrows=1, ncols=1, figsize=None, imsize=3, add_vert=0, **kwargs):
    if figsize is None: figsize=(ncols*imsize, nrows*imsize+add_vert)
    fig,ax = plt.subplots(nrows, ncols, figsize=figsize, **kwargs)
    if nrows*ncols==1: ax = array([ax])
    return fig,ax

# Cell
def _fig_bounds(x):
    r = x//32
    return min(5, max(1,r))

# Cell
def show_image(im, ax=None, figsize=None, title=None, ctx=None, **kwargs):
    "Show a PIL or PyTorch image on `ax`."
    # Handle pytorch axis order
    if hasattrs(im, ('data','cpu','permute')):
        im = im.data.cpu()
        if im.shape[0]<5: im=im.permute(1,2,0)
    elif not isinstance(im,np.ndarray): im=array(im)
    # Handle 1-channel images
    if im.shape[-1]==1: im=im[...,0]

    ax = ifnone(ax,ctx)
    if figsize is None: figsize = (_fig_bounds(im.shape[0]), _fig_bounds(im.shape[1]))
    if ax is None: _,ax = plt.subplots(figsize=figsize)
    ax.imshow(im, **kwargs)
    if title is not None: ax.set_title(title)
    ax.axis('off')
    return ax

# Cell
def show_titled_image(o, **kwargs):
    "Call `show_image` destructuring `o` to `(img,title)`"
    show_image(o[0], title=str(o[1]), **kwargs)

# Cell
@delegates(subplots)
def show_images(ims, nrows=1, ncols=None, titles=None, **kwargs):
    "Show all images `ims` as subplots with `rows` using `titles`"
    if ncols is None: ncols = int(math.ceil(len(ims)/nrows))
    if titles is None: titles = [None]*len(ims)
    axs = subplots(nrows, ncols, **kwargs)[1].flat
    for im,t,ax in zip(ims, titles, axs): show_image(im, ax=ax, title=t)

# Cell
class ArrayBase(ndarray):
    @classmethod
    def _before_cast(cls, x): return x if isinstance(x,ndarray) else array(x)

# Cell
class ArrayImageBase(ArrayBase):
    _show_args = {'cmap':'viridis'}
    def show(self, ctx=None, **kwargs):
        return show_image(self, ctx=ctx, **{**self._show_args, **kwargs})

# Cell
class ArrayImage(ArrayImageBase): pass

# Cell
class ArrayImageBW(ArrayImage): _show_args = {'cmap':'Greys'}

# Cell
class ArrayMask(ArrayImageBase): _show_args = {'alpha':0.5, 'cmap':'tab20', 'interpolation':'nearest'}

# Cell
@patch
def __array_eq__(self:Tensor,b):
    return torch.equal(self,b) if self.dim() else self==b

# Cell
def _array2tensor(x):
    if x.dtype==np.uint16: x = x.astype(np.float32)
    return torch.from_numpy(x)

# Cell
def tensor(x, *rest, **kwargs):
    "Like `torch.as_tensor`, but handle lists too, and can pass multiple vector elements directly."
    if len(rest): x = (x,)+rest
    # There was a Pytorch bug in dataloader using num_workers>0. Haven't confirmed if fixed
    # if isinstance(x, (tuple,list)) and len(x)==0: return tensor(0)
    res = (x if isinstance(x, Tensor)
           else torch.tensor(x, **kwargs) if isinstance(x, (tuple,list))
           else _array2tensor(x) if isinstance(x, ndarray)
           else as_tensor(x.values, **kwargs) if isinstance(x, (pd.Series, pd.DataFrame))
           else as_tensor(x, **kwargs) if hasattr(x, '__array__') or is_iter(x)
           else _array2tensor(array(x), **kwargs))
    if res.dtype is torch.float64: return res.float()
    return res

# Cell
def set_seed(s):
    "Set random seed for `random`, `torch`, and `numpy` (where available)"
    try: torch.manual_seed(s)
    except NameError: pass
    try: np.random.seed(s%(2**32-1))
    except NameError: pass
    random.seed(s)

# Cell
def unsqueeze(x, dim=-1, n=1):
    "Same as `torch.unsqueeze` but can add `n` dims"
    for _ in range(n): x = x.unsqueeze(dim)
    return x

# Cell
def unsqueeze_(x, dim=-1, n=1):
    "Same as `torch.unsqueeze_` but can add `n` dims"
    for _ in range(n): x.unsqueeze_(dim)
    return x

# Cell
def _fa_rebuild_tensor (cls, *args, **kwargs): return cls(torch._utils._rebuild_tensor_v2(*args, **kwargs))
def _fa_rebuild_qtensor(cls, *args, **kwargs): return cls(torch._utils._rebuild_qtensor  (*args, **kwargs))

# Cell
def apply(func, x, *args, **kwargs):
    "Apply `func` recursively to `x`, passing on args"
    if is_listy(x): return type(x)([apply(func, o, *args, **kwargs) for o in x])
    if isinstance(x,dict):  return {k: apply(func, v, *args, **kwargs) for k,v in x.items()}
    res = func(x, *args, **kwargs)
    return res if x is None else retain_type(res, x)

# Cell
def maybe_gather(x, axis=0):
    "Gather copies of `x` on `axis` (if training is distributed)"
    if num_distrib()<=1: return x
    ndim = x.ndim
    res = [x.new_zeros(*x.shape if ndim > 0 else (1,)) for _ in range(num_distrib())]
    torch.distributed.all_gather(res, x if ndim > 0 else x[None])
    return torch.cat(res, dim=axis) if ndim > 0 else torch.cat(res, dim=axis).mean()

# Cell
def to_detach(b, cpu=True, gather=True):
    "Recursively detach lists of tensors in `b `; put them on the CPU if `cpu=True`."
    def _inner(x, cpu=True, gather=True):
        if not isinstance(x,Tensor): return x
        x = x.detach()
        if gather: x = maybe_gather(x)
        return x.cpu() if cpu else x
    return apply(_inner, b, cpu=cpu, gather=gather)

# Cell
def to_half(b):
    "Recursively map lists of tensors in `b ` to FP16."
    return apply(lambda x: x.half() if torch.is_floating_point(x) else x, b)

# Cell
def to_float(b):
    "Recursively map lists of int tensors in `b ` to float."
    return apply(lambda x: x.float() if torch.is_floating_point(x) else x, b)

# Cell
# None: True if available; True: error if not availabe; False: use CPU
defaults.use_cuda = None

# Cell
def default_device(use_cuda=-1):
    "Return or set default device; `use_cuda`: None - CUDA if available; True - error if not availabe; False - CPU"
    if use_cuda != -1: defaults.use_cuda=use_cuda
    use = defaults.use_cuda or (torch.cuda.is_available() and defaults.use_cuda is None)
    assert torch.cuda.is_available() or not use
    return torch.device(torch.cuda.current_device()) if use else torch.device('cpu')

# Cell
def to_device(b, device=None):
    "Recursively put `b` on `device`."
    if defaults.use_cuda==False: device='cpu'
    elif device is None: device=default_device()
    def _inner(o): return o.to(device, non_blocking=True) if isinstance(o,Tensor) else o.to_device(device) if hasattr(o, "to_device") else o
    return apply(_inner, b)

# Cell
def to_cpu(b):
    "Recursively map lists of tensors in `b ` to the cpu."
    return to_device(b,'cpu')

# Cell
def to_np(x):
    "Convert a tensor to a numpy array."
    return apply(lambda o: o.data.cpu().numpy(), x)

# Cell
def to_concat(xs, dim=0):
    "Concat the element in `xs` (recursively if they are tuples/lists of tensors)"
    if is_listy(xs[0]): return type(xs[0])([to_concat([x[i] for x in xs], dim=dim) for i in range_of(xs[0])])
    if isinstance(xs[0],dict):  return {k: to_concat([x[k] for x in xs], dim=dim) for k in xs[0].keys()}
    #We may receives xs that are not concatenatable (inputs of a text classifier for instance),
    #   in this case we return a big list
    try:    return retain_type(torch.cat(xs, dim=dim), xs[0])
    except: return sum([L(retain_type(o_.index_select(dim, tensor(i)).squeeze(dim), xs[0])
                          for i in range_of(o_)) for o_ in xs], L())

# Cell
@patch
def set_meta(self:Tensor, x):
    "Set all metadata in `__dict__`"
    if hasattr(x,'__dict__'): self.__dict__ = x.__dict__

# Cell
@patch
def get_meta(self:Tensor, n, d=None):
    "Set `n` from `self._meta` if it exists and returns default `d` otherwise"
    return getattr(self, '_meta', {}).get(n, d)

# Cell
@patch
def as_subclass(self:Tensor, typ):
    "Cast to `typ` (should be in future PyTorch version, so remove this then)"
    res = torch.Tensor._make_subclass(typ, self)
    return retain_meta(self, res)

# Cell
class TensorBase(Tensor):
    def __new__(cls, x, **kwargs):
        global printed_device2
        res = cast(tensor(x), cls)
        if res.device != tpu_device:
          print(f"************** XXX ######################### EL DEVICE ES {tpu_device} no es {res.device} ni el de entrada {x.device} ######################### **************") if printed_device < 10 else noop
          printed_device2 += 1
        res._meta = kwargs
        return res

    @classmethod
    def _before_cast(cls, x): return x if isinstance(x,Tensor) else tensor(x)

    def __reduce_ex__(self,proto):
        torch.utils.hooks.warn_if_has_hooks(self)
        args = (type(self), self.storage(), self.storage_offset(), tuple(self.size()), self.stride())
        if self.is_quantized: args = args + (self.q_scale(), self.q_zero_point())
        f = _fa_rebuild_qtensor if self.is_quantized else  _fa_rebuild_tensor
        return (f, args + (self.requires_grad, OrderedDict()))

    def gi(self, i):
        res = self[i]
        return res.as_subclass(type(self)) if isinstance(res,Tensor) else res

    def __repr__(self):
        return re.sub('tensor', self.__class__.__name__, super().__repr__())

# Cell
def _patch_tb():
    # TRACE: print("*********################### ######################## patch_tb")
    if getattr(TensorBase,'_patched',False): return
    TensorBase._patched = True

    def get_f(fn):
        def _f(self, *args, **kwargs):
            cls = self.__class__
            res = getattr(super(TensorBase, self), fn)(*args, **kwargs)
            return retain_type(res, self)
        return _f

    t = tensor([1])
    skips = 'as_subclass imag real __getitem__ __class__ __deepcopy__ __delattr__ __dir__ __doc__ __getattribute__ __hash__ __init__ \
        __init_subclass__ __new__ __reduce__ __reduce_ex__ __repr__ __module__ __setstate__'.split()

    for fn in dir(t):
        if fn in skips: continue
        f = getattr(t, fn)
        if isinstance(f, (MethodWrapperType, BuiltinFunctionType, BuiltinMethodType, MethodType, FunctionType)):
            setattr(TensorBase, fn, get_f(fn))

_patch_tb()

# Cell
class TensorCategory(TensorBase): pass

# Cell
class TensorMultiCategory(TensorCategory): pass

# Cell
class TensorImageBase(TensorBase):
    _show_args = ArrayImageBase._show_args
    def show(self, ctx=None, **kwargs):
        return show_image(self, ctx=ctx, **{**self._show_args, **kwargs})

# Cell
class TensorImage(TensorImageBase): pass

# Cell
class TensorImageBW(TensorImage): _show_args = ArrayImageBW._show_args

# Cell
class TensorMask(TensorImageBase):
    _show_args = ArrayMask._show_args

    def show(self, ctx=None, **kwargs):
        codes = self.get_meta('codes')
        if codes is not None: kwargs = merge({'vmin': 1, 'vmax': len(codes)}, kwargs)
        return super().show(ctx=ctx, **kwargs)

# Cell
class TitledTensorScalar(TensorBase):
    "A tensor containing a scalar that has a `show` method"
    def show(self, **kwargs): show_title(self.item(), **kwargs)

# Cell
@patch
def tensored(self:L):
    "`mapped(tensor)`"
    return self.map(tensor)
@patch
def stack(self:L, dim=0):
    "Same as `torch.stack`"
    return torch.stack(list(self.tensored()), dim=dim)
@patch
def cat  (self:L, dim=0):
    "Same as `torch.cat`"
    return torch.cat  (list(self.tensored()), dim=dim)

# Cell
def concat(*ls):
    "Concatenate tensors, arrays, lists, or tuples"
    if not len(ls): return []
    it = ls[0]
    if isinstance(it,torch.Tensor): res = torch.cat(ls)
    elif isinstance(it,ndarray): res = np.concatenate(ls)
    else:
        res = itertools.chain.from_iterable(map(L,ls))
        if isinstance(it,(tuple,list)): res = type(it)(res)
        else: res = L(res)
    return retain_type(res, it)

# Cell
class Chunks:
    "Slice and int indexing into a list of lists"
    def __init__(self, chunks, lens=None):
        self.chunks = chunks
        self.lens = L(map(len,self.chunks) if lens is None else lens)
        self.cumlens = np.cumsum(0+self.lens)
        self.totlen = self.cumlens[-1]

    def __getitem__(self,i):
        if isinstance(i,slice): return retain_type(self.getslice(i), old=self.chunks[0])
        di,idx = self.doc_idx(i)
        return retain_type(self.chunks[di][idx], old=self.chunks[0])

    def getslice(self, i):
        st_d,st_i = self.doc_idx(ifnone(i.start,0))
        en_d,en_i = self.doc_idx(ifnone(i.stop,self.totlen+1))
        res = [self.chunks[st_d][st_i:(en_i if st_d==en_d else sys.maxsize)]]
        for b in range(st_d+1,en_d): res.append(self.chunks[b])
        if st_d!=en_d and en_d<len(self.chunks): res.append(self.chunks[en_d][:en_i])
        return concat(*res)

    def doc_idx(self, i):
        if i<0: i=self.totlen+i # count from end
        docidx = np.searchsorted(self.cumlens, i+1)-1
        cl = self.cumlens[docidx]
        return docidx,i-cl

# Cell
def show_title(o, ax=None, ctx=None, label=None, color='black', **kwargs):
    "Set title of `ax` to `o`, or print `o` if `ax` is `None`"
    ax = ifnone(ax,ctx)
    if ax is None: print(o)
    elif hasattr(ax, 'set_title'):
        t = ax.title.get_text()
        if len(t) > 0: o = t+'\n'+str(o)
        ax.set_title(o, color=color)
    elif isinstance(ax, pd.Series):
        while label in ax: label += '_'
        ax = ax.append(pd.Series({label: o}))
    return ax

# Cell
class ShowTitle:
    "Base class that adds a simple `show`"
    _show_args = {'label': 'text'}
    def show(self, ctx=None, **kwargs):
        "Show self"
        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))

class TitledInt(Int, ShowTitle):
    _show_args = {'label': 'text'}
    def show(self, ctx=None, **kwargs):
        "Show self"
        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))

class TitledFloat(Float, ShowTitle):
    _show_args = {'label': 'text'}
    def show(self, ctx=None, **kwargs):
        "Show self"
        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))

class TitledStr(Str, ShowTitle):
    _show_args = {'label': 'text'}
    def show(self, ctx=None, **kwargs):
        "Show self"
        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))

class TitledTuple(Tuple, ShowTitle):
    _show_args = {'label': 'text'}
    def show(self, ctx=None, **kwargs):
        "Show self"
        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))

add_docs(TitledInt, "An `int` with `show`"); add_docs(TitledStr, "An `str` with `show`");
add_docs(TitledFloat, "A `float` with `show`"); add_docs(TitledTuple, "A `Tuple` with `show`")

# Cell
@patch
def truncate(self:TitledStr, n):
    "Truncate self to `n`"
    words = self.split(' ')[:n]
    return TitledStr(' '.join(words))

# Cell
if not hasattr(pd.DataFrame,'_old_init'): pd.DataFrame._old_init = pd.DataFrame.__init__

# Cell
@patch
def __init__(self:pd.DataFrame, data=None, index=None, columns=None, dtype=None, copy=False):
    if data is not None and isinstance(data, Tensor): data = to_np(data)
    self._old_init(data, index=index, columns=columns, dtype=dtype, copy=copy)

# Cell
def get_empty_df(n):
    "Return `n` empty rows of a dataframe"
    df = pd.DataFrame(index = range(n))
    return [df.iloc[i] for i in range(n)]

# Cell
def display_df(df):
    "Display `df` in a notebook or defaults to print"
    try: from IPython.display import display, HTML
    except: return print(df)
    display(HTML(df.to_html()))

# Cell
def get_first(c):
    "Get the first element of c, even if c is a dataframe"
    return getattr(c, 'iloc', c)[0]

# Cell
def one_param(m):
    "First parameter in `m`"
    return first(m.parameters())

# Cell
def item_find(x, idx=0):
    "Recursively takes the `idx`-th element of `x`"
    if is_listy(x): return item_find(x[idx])
    if isinstance(x,dict):
        key = list(x.keys())[idx] if isinstance(idx, int) else idx
        return item_find(x[key])
    return x

# Cell
def find_device(b):
    "Recursively search the device of `b`."
    return item_find(b).device

# Cell
def find_bs(b):
    "Recursively search the batch size of `b`."
    return item_find(b).shape[0]

# Cell
class Module(nn.Module, metaclass=PrePostInitMeta):
    "Same as `nn.Module`, but no need for subclasses to call `super().__init__`"
    def __pre_init__(self, *args, **kwargs): super().__init__()
    def __init__(self): pass

# Cell
from torch.nn.parallel import DistributedDataParallel

def get_model(model):
    "Return the model maybe wrapped inside `model`."
    return model.module if isinstance(model, (DistributedDataParallel, nn.DataParallel)) else model

# Cell
def one_hot(x, c):
    "One-hot encode `x` with `c` classes."
    res = torch.zeros(c, dtype=torch.uint8)
    if isinstance(x, Tensor) and x.numel()>0: res[x] = 1.
    else: res[list(L(x, use_list=None))] = 1.
    return res

# Cell
def one_hot_decode(x, vocab=None):
    return L(vocab[i] if vocab else i for i,x_ in enumerate(x) if x_==1)

# Cell
def params(m):
    "Return all parameters of `m`"
    return [p for p in m.parameters()]

# Cell
def trainable_params(m):
    "Return all trainable parameters of `m`"
    return [p for p in m.parameters() if p.requires_grad]

# Cell
norm_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d)

# Cell
def bn_bias_params(m, with_bias=True): # TODO: Rename to `norm_bias_params`
    "Return all bias and BatchNorm parameters"
    if isinstance(m, norm_types): return L(m.parameters())
    res = L(m.children()).map(bn_bias_params, with_bias=with_bias).concat()
    if with_bias and getattr(m, 'bias', None) is not None: res.append(m.bias)
    return res

# Cell
def batch_to_samples(b, max_n=10):
    "'Transposes' a batch to (at most `max_n`) samples"
    if isinstance(b, Tensor): return retain_types(list(b[:max_n]), [b])
    else:
        res = L(b).map(partial(batch_to_samples,max_n=max_n))
        return retain_types(res.zip(), [b])

# Cell
@patch
def interp_1d(x:Tensor, xp, fp):
    "Same as `np.interp`"
    slopes = (fp[1:]-fp[:-1])/(xp[1:]-xp[:-1])
    incx = fp[:-1] - (slopes*xp[:-1])
    locs = (x[:,None]>=xp[None,:]).long().sum(1)-1
    locs = locs.clamp(0,len(slopes)-1)
    return slopes[locs]*x + incx[locs]

# Cell
@patch
def pca(x:Tensor, k=2):
    "Compute PCA of `x` with `k` dimensions."
    x = x-torch.mean(x,0)
    U,S,V = torch.svd(x.t())
    return torch.mm(x,U[:,:k])

# Cell
def logit(x):
    "Logit of `x`, clamped to avoid inf."
    x = x.clamp(1e-7, 1-1e-7)
    return -(1/x-1).log()

# Cell
def num_distrib():
    "Return the number of processes in distributed training (if applicable)."
    return int(os.environ.get('WORLD_SIZE', 0))

# Cell
def rank_distrib():
    "Return the distributed rank of this process (if applicable)."
    return int(os.environ.get('RANK', 0))

# Cell
def distrib_barrier():
    "Place a synchronization barrier in distributed training so that ALL sub-processes in the pytorch process group must arrive here before proceeding."
    if num_distrib() > 1 and torch.distributed.is_initialized(): torch.distributed.barrier()

# Cell
# Saving arrays requires pytables - optional dependency
try: import tables
except: pass

# Cell
def _comp_filter(lib='lz4',lvl=3): return tables.Filters(complib=f'blosc:{lib}', complevel=lvl)

# Cell
@patch
def save_array(p:Path, o, complib='lz4', lvl=3):
    "Save numpy array to a compressed `pytables` file, using compression level `lvl`"
    if isinstance(o,Tensor): o = to_np(o)
    with tables.open_file(p, mode='w', filters=_comp_filter(lib=complib,lvl=lvl)) as f: f.create_carray('/', 'data', obj=o)

# Cell
@patch
def load_array(p:Path):
    "Save numpy array to a `pytables` file"
    with tables.open_file(p, 'r') as f: return f.root.data.read()

# Cell
def base_doc(elt):
    "Print a base documentation of `elt`"
    name = getattr(elt, '__qualname__', getattr(elt, '__name__', ''))
    print(f'{name}{inspect.signature(elt)}\n{inspect.getdoc(elt)}\n')
    print('To get a prettier result with hyperlinks to source code and documentation, install nbdev: pip install nbdev')

# Cell
def doc(elt):
    "Try to use doc form nbdev and fall back to `base_doc`"
    try:
        from nbdev.showdoc import doc
        doc(elt)
    except: base_doc(elt)

# Cell
def nested_reorder(t, idxs):
    "Reorder all tensors in `t` using `idxs`"
    if isinstance(t, (Tensor,L)): return t[idxs]
    elif is_listy(t): return type(t)(nested_reorder(t_, idxs) for t_ in t)
    if t is None: return t
    raise TypeError(f"Expected tensor, tuple, list or L but got {type(t)}")

# Cell
def to_image(x):
    if isinstance(x,Image.Image): return x
    if isinstance(x,Tensor): x = to_np(x.permute((1,2,0)))
    if x.dtype==np.float32: x = (x*255).astype(np.uint8)
    return Image.fromarray(x, mode=['RGB','CMYK'][x.shape[0]==4])

# Cell
def make_cross_image(bw=True):
    "Create a tensor containing a cross image, either `bw` (True) or color"
    if bw:
        im = torch.zeros(5,5)
        im[2,:] = 1.
        im[:,2] = 1.
    else:
        im = torch.zeros(3,5,5)
        im[0,2,:] = 1.
        im[1,:,2] = 1.
    return im

# Cell
def show_image_batch(b, show=show_titled_image, items=9, cols=3, figsize=None, **kwargs):
    "Display batch `b` in a grid of size `items` with `cols` width"
    if items<cols: cols=items
    rows = (items+cols-1) // cols
    if figsize is None: figsize = (cols*3, rows*3)
    fig,axs = plt.subplots(rows, cols, figsize=figsize)
    for *o,ax in zip(*to_cpu(b), axs.flatten()): show(o, ax=ax, **kwargs)

# Cell
def requires_grad(m):
    "Check if the first parameter of `m` requires grad or not"
    ps = list(m.parameters())
    return ps[0].requires_grad if len(ps)>0 else False

# Cell
def init_default(m, func=nn.init.kaiming_normal_):
    "Initialize `m` weights with `func` and set `bias` to 0."
    if func:
        if hasattr(m, 'weight'): func(m.weight)
        if hasattr(m, 'bias') and hasattr(m.bias, 'data'): m.bias.data.fill_(0.)
    return m

# Cell
def cond_init(m, func):
    "Apply `init_default` to `m` unless it's a batchnorm module"
    if (not isinstance(m, norm_types)) and requires_grad(m): init_default(m, func)

# Cell
def apply_leaf(m, f):
    "Apply `f` to children of `m`."
    c = m.children()
    if isinstance(m, nn.Module): f(m)
    for l in c: apply_leaf(l,f)

# Cell
def apply_init(m, func=nn.init.kaiming_normal_):
    "Initialize all non-batchnorm layers of `m` with `func`."
    apply_leaf(m, partial(cond_init, func=func))

# Cell
from multiprocessing import Process, Queue

# Cell
def set_num_threads(nt):
    "Get numpy (and others) to use `nt` threads"
    try: import mkl; mkl.set_num_threads(nt)
    except: pass
    torch.set_num_threads(1)
    os.environ['IPC_ENABLE']='1'
    for o in ['OPENBLAS_NUM_THREADS','NUMEXPR_NUM_THREADS','OMP_NUM_THREADS','MKL_NUM_THREADS']:
        os.environ[o] = str(nt)

# Cell
@delegates(concurrent.futures.ProcessPoolExecutor)
class ProcessPoolExecutor(concurrent.futures.ProcessPoolExecutor):
    def __init__(self, max_workers=None, on_exc=print, **kwargs):
        self.not_parallel = max_workers==0
        self.on_exc = on_exc
        if self.not_parallel: max_workers=1
        super().__init__(max_workers, **kwargs)

    def map(self, f, items, *args, **kwargs):
        g = partial(f, *args, **kwargs)
        if self.not_parallel: return map(g, items)
        try: return super().map(g, items)
        except Exception as e: self.on_exc(e)

# Cell
def parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=True, **kwargs):
    "Applies `func` in parallel to `items`, using `n_workers`"
    with ProcessPoolExecutor(n_workers) as ex:
        r = ex.map(f,items, *args, **kwargs)
        if progress:
            if total is None: total = len(items)
            r = progress_bar(r, total=total, leave=False)
        return L(r)

# Cell
def run_procs(f, f_done, args):
    "Call `f` for each item in `args` in parallel, yielding `f_done`"
    processes = L(args).map(Process, args=arg0, target=f)
    for o in processes: o.start()
    try: yield from f_done()
    except Exception as e: print(e)
    finally: processes.map(Self.join())

# Cell
def parallel_gen(cls, items, n_workers=defaults.cpus, as_gen=False, **kwargs):
    "Instantiate `cls` in `n_workers` procs & call each on a subset of `items` in parallel."
    batches = np.array_split(items, n_workers)
    idx = np.cumsum(0 + L(batches).map(len))
    queue = Queue()
    def f(batch, start_idx):
        for i,b in enumerate(cls(**kwargs)(batch)): queue.put((start_idx+i,b))
    def done(): return (queue.get() for _ in progress_bar(items, leave=False))
    yield from run_procs(f, done, L(batches,idx).zip())

# Cell
def script_use_ctx(f):
    "Decorator: create jit script and pass everything in `ctx.saved_variables to `f`, after `*args`"
    sf = torch.jit.script(f)
    def _f(ctx, *args, **kwargs): return sf(*args, *ctx.saved_variables, **kwargs)
    return update_wrapper(_f,f)

# Cell
def script_save_ctx(static, *argidx):
    "Decorator: create jit script and save args with indices `argidx` using `ctx.save_for_backward`"
    def _dec(f):
        sf = torch.jit.script(f)
        def _f(ctx, *args, **kwargs):
            if argidx:
                save = [args[o] for o in argidx]
                ctx.save_for_backward(*save)
            if not argidx: args = [ctx]+args
            return sf(*args, **kwargs)
        if static: _f = staticmethod(_f)
        return update_wrapper(_f,f)
    return _dec

# Cell
def script_fwd(*argidx):
    "Decorator: create static jit script and save args with indices `argidx` using `ctx.save_for_backward`"
    return script_save_ctx(True, *argidx)

# Cell
def script_bwd(f):
    "Decorator: create static jit script and pass everything in `ctx.saved_variables to `f`, after `*args`"
    return staticmethod(script_use_ctx(f))

# Cell
def grad_module(cls):
    "Decorator: convert `cls` into an autograd function"
    class _c(nn.Module):
        def forward(self, *args, **kwargs): return cls.apply(*args, **kwargs)
    return _c

# Comes from 13b_metrics.ipynb, cell
def flatten_check(inp, targ):
    "Check that `out` and `targ` have the same number of elements and flatten them."
    inp,targ = inp.contiguous().view(-1),targ.contiguous().view(-1)
    test_eq(len(inp), len(targ))
    return inp,targ

print("---------------------- PASSED torch_base.py")

print("---------------------------------------------------------- ALMOOOOST END")

################################################################33

from datetime import datetime

@patch_to(Learner)
def one_batch(self, i, b):
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                   Learner#one_batch@BATCH {i}")
        self.iter = i
        try:
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                   Learner#one_batch@SPLIT {i}")
            self._split(b);                                  self('begin_batch')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called begin_batch")
            self.pred = self.model(*self.xb);                self('after_pred')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called after_pred")
            if len(self.yb) == 0: return
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} did not return becasue yb.length == 0")
            self.loss = self.loss_func(self.pred, *self.yb); self('after_loss')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called after_loss")
            if not self.training: return
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} did not return because we are training")
            self.loss.backward();                            self('after_backward')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called after_backward")
            #self.opt.step();                                 self('after_step')
            xm.optimizer_step(self.opt,barrier=True);        self('after_step')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called after_step")
            self.opt.zero_grad()
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called zero_grad!!!!!!!")
        except CancelBatchException:
            self('after_cancel_batch')
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 Learner#one_batch@ {i} called after_cancel_batch")
        finally:
            if False:
                # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} AFTER BATCH NOT CALLED")
                # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} AFTER BATCH NOT CALLED")
                # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} AFTER BATCH NOT CALLED")
                # TRACE: print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} AFTER BATCH NOT CALLED")
                print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} AFTER BATCH NOT CALLED")
            else:
                print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} WILL CALL after_batch")
                self('after_batch')
                print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}                 {i} called after_batch")

from torch.utils.data.dataloader import _MultiProcessingDataLoaderIter,_SingleProcessDataLoaderIter,_DatasetKind
_loaders = (_MultiProcessingDataLoaderIter,_SingleProcessDataLoaderIter)

@patch_to(DataLoader)
def __iter__(self):
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#__iter__                         0")
        self.randomize()
        self.before_iter()
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#__iter__ START FOR               1")
        xxxx = _loaders[self.fake_l.num_workers==0](self.fake_l)
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#__iter__ ----------------------- 1")
        for b in xxxx:
            print("======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
            print("======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
            print("======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------")
            if self.device is not None:
                print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#iterator to device from {b[0].device} y {b[1].device} to {self.device}")
                b = to_device(b, self.device)
                print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#iterator to done!!!!")
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch({b[0].device} y {b[1].device}) len of b is {len(b)}")
            yield self.after_batch(b)
            print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#yielding                    4!!!!")
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}  DataLoader#DataLoader#DataLoader#__iter__ END FOR                 2")
        self.after_iter()
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}   DataLoader#DataLoader#DataLoader#__iter__ after ITER")
        if hasattr(self, 'it'): delattr(self, 'it')
        print(f"{datetime.now().strftime(' (%H:%M:%S.%f)')}     DataLoader#DataLoader#DataLoader#END __iter__")

@patch_to(Learner)
def all_batches(self):
        print("                                          Learner#ALL_BATCHES")
        self.n_iter = len(self.dl)
        print("                                          Learner#ALL_BATCHES GEN enumerator!!!!")
        e = enumerate(self.dl)
        print("                                          Learner#ALL_BATCHES generate enumerator!!!!")
        for o in e:
            print(">>>> =======================================================================")
            print(">>>> =======================================================================")
            print(f">>>>                                           Learner#ALL_BATCHES CALL ENTER [{o[0]}]")
            print(">>>> =======================================================================")
            print(">>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch")
            self.one_batch(*o)
            print("<<<< ======================================================================= OOOOOOOOOOOOOOOO")
            print("<<<< =======================================================================")
            print("<<<< =======================================================================")
            print("<<<<                                           Learner#ALL_BATCHES CALL EXIT")
            print("<<<< ====================================================================***")
            print("<<<< ========================================*******************************")
            print("<<<< ===========================********************************************")

@patch_to(Learner)
def create_opt(self):
        # TRACE: print_local("create_opt!!!")
        ooo = self.opt_func(self.splitter(self.model), lr=self.lr)
        prox = XLAOptimProxy(ooo)
        self.opt = prox
        if not self.wd_bn_bias:
            for p in self._bn_bias_state(True ): p['do_wd'] = False
        if self.train_bn:
            for p in self._bn_bias_state(False): p['force_train'] = True

#proxyLearn = Learner(dls_tpu, Lenet2(), metrics=accuracy, opt_func=Adam)#, cbs=CallbackXLA)

#print("---------------------------------------------------------- Before fit")
#proxyLearn.fit(1, 10e-3) # 0.05) NOTE: Im not sure if this works...!!! it is now 96!
#print("---------------------------------------------------------- After fit")

print(".............................................\n"*20)

@patch
@log_args(but_as=Learner.fit)
def fit_one_cycle(self:Learner, n_epoch, lr_max=None, div=25., div_final=1e5, pct_start=0.25, wd=None,
                  moms=None, cbs=None, reset_opt=False):
    "Fit `self.model` for `n_epoch` using the 1cycle policy."
    if self.opt is None: self.create_opt()
    self.opt.set_hyper('lr', self.lr if lr_max is None else lr_max)
    lr_max = np.array([h['lr'] for h in self.opt.hypers])
    scheds = {'lr': combined_cos(pct_start, lr_max/div, lr_max, lr_max/div_final),
              'mom': combined_cos(pct_start, *(self.moms if moms is None else moms))}
    self.fit(n_epoch, cbs=ParamScheduler(scheds)+L(cbs), reset_opt=reset_opt, wd=wd)

@patch
@log_args(but_as=Learner.fit)
@delegates(Learner.fit_one_cycle)
def fine_tune(self:Learner, epochs, base_lr=2e-3, freeze_epochs=1, lr_mult=100,
              pct_start=0.3, div=5.0, **kwargs):
    "Fine tune with `freeze` for `freeze_epochs` then with `unfreeze` from `epochs` using discriminative LR"
    self.freeze()
    self.fit_one_cycle(freeze_epochs, slice(base_lr), pct_start=0.99, **kwargs)
    base_lr /= 2
    self.unfreeze()
    self.fit_one_cycle(epochs, slice(base_lr/lr_mult, base_lr), pct_start=pct_start, div=div, **kwargs)

@patch
def freeze_to(self:Learner, n):
    if self.opt is None: self.create_opt()
    self.opt.freeze_to(n)
    self.opt.clear_state()

@patch
def freeze(self:Learner): self.freeze_to(-1)

@patch
def unfreeze(self:Learner): self.freeze_to(0)

def _add_norm(dls, meta, pretrained):
    if not pretrained: return
    after_batch = dls.after_batch
    if first(o for o in after_batch.fs if isinstance(o,Normalize)): return
    stats = meta.get('stats')
    if stats is None: return
    after_batch.add(Normalize.from_stats(*stats))

def default_split(m):
    "Default split of a model between body and head"
    return L(m[0], m[1:]).map(params)
_default_meta    = {'cut':None, 'split':default_split}

@log_args(to_return=True, but_as=Learner.__init__)
@delegates(Learner.__init__)
def cnn_learner(dls, arch, loss_func=None, pretrained=True, cut=None, splitter=None,
                y_range=None, config=None, n_out=None, normalize=True, **kwargs):
    "Build a convnet style learner from `dls` and `arch`"
    print(f"cnn dls.device {dls.device}")
    if config is None: config = {}
    meta = model_meta.get(arch, _default_meta)
    if n_out is None: n_out = get_c(dls)
    assert n_out, "`n_out` is not defined, and could not be infered from data, set `dls.c` or pass `n_out`"
    if normalize: _add_norm(dls, meta, pretrained)
    if y_range is None and 'y_range' in config: y_range = config.pop('y_range')
    print(f"cnn dls.device MIDDLE")
    model = create_cnn_model(arch, n_out, ifnone(cut, meta['cut']), pretrained, y_range=y_range, **config)
    learn = Learner(dls, model, loss_func=loss_func, splitter=ifnone(splitter, meta['split']), **kwargs)
    if pretrained: learn.freeze()
    print(f"cnn dls.device END")
    return learn

path = untar_data(URLs.PETS)/'images'
pat = r'(.+)_\d+.jpg$'
datablock = DataBlock(
    blocks=(ImageBlock,CategoryBlock),
    get_items=get_image_files,
    splitter=RandomSplitter(seed=42),
    get_y=using_attr(RegexLabeller(pat),'name'),
    item_tfms=Resize(224),
    # batch_tfms=[]
    batch_tfms=aug_transforms(size=224,min_scale=0.75)
)
datablock.summary(path)
dls = datablock.dataloaders(path,bs=256, device=tpu_device)
print("CNN Learner")
learner = cnn_learner(dls, resnet34, metrics=accuracy)
print("FINE TUNE")
learner.fine_tune(1,base_lr=4e-3,freeze_epochs=2)
print("end FINE TUNE")
```

    ---------------------------------------------------------- START
    device is xla:1
    ---------------------------------------------------------- START:1
    ---------------------------------------------------------- START:2
    ---------------------------------------------------------- START:3
     (07:15:02.164228)  DataLoader#DataLoader#DataLoader#__iter__                         0
     (07:15:02.165292)  DataLoader#DataLoader#DataLoader#__iter__ START FOR               1
     (07:15:02.169795)  DataLoader#DataLoader#DataLoader#__iter__ ----------------------- 1
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:15:02.193488)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:15:02.193909)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:15:02.201670)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:15:02.202112)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    ---------------------------------------------------------- START:4
    ---------------------------------------------------------- START:5
    ---------------------------------------------------------- START:6
    ---------------------------------------------------------- START:7
    ---------------------------------------------------------- START:8 enbd STARET
    ---------------------- STARTED class Learner
    ---------------------- PASSED LEARNER
    ---------------------- STARTED torch_base.py
    ---------------------- PASSED torch_base.py
    ---------------------------------------------------------- ALMOOOOST END
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................
    .............................................

    Setting-up type transforms pipelines
    Collecting items from /root/.fastai/data/oxford-iiit-pet/images
    Found 7390 items
    2 datasets of sizes 5912,1478
    Setting up Pipeline: PILBase.create
    Setting up Pipeline: partial -> Categorize

    Building one sample
      Pipeline: PILBase.create
        starting from
          /root/.fastai/data/oxford-iiit-pet/images/scottish_terrier_127.jpg
        applying PILBase.create gives
          PILImage mode=RGB size=375x500
      Pipeline: partial -> Categorize
        starting from
          /root/.fastai/data/oxford-iiit-pet/images/scottish_terrier_127.jpg
        applying partial gives
          scottish_terrier
        applying Categorize gives
          TensorCategory(32)

    Final sample: (PILImage mode=RGB size=375x500, TensorCategory(32))

    Setting up after_item: Pipeline: Resize -> ToTensor
    Setting up before_batch: Pipeline:
    Setting up after_batch: Pipeline: IntToFloatTensor -> AffineCoordTfm -> RandomResizedCropGPU -> LightingTfm

    Building one batch
    Applying item_tfms to the first sample:
      Pipeline: Resize -> ToTensor
        starting from
          (PILImage mode=RGB size=375x500, TensorCategory(32))
        applying Resize gives
          (PILImage mode=RGB size=224x224, TensorCategory(32))
        applying ToTensor gives
          (TensorImage of size 3x224x224, TensorCategory(32))

    Adding the next 3 samples

    No before_batch transform to apply

    Collating items in a batch

    Applying batch_tfms to the batch built
      Pipeline: IntToFloatTensor -> AffineCoordTfm -> RandomResizedCropGPU -> LightingTfm
        starting from
          (TensorImage of size 4x3x224x224, TensorCategory([32, 20, 29, 35]))
        applying IntToFloatTensor gives
          (TensorImage of size 4x3x224x224, TensorCategory([32, 20, 29, 35]))
        applying AffineCoordTfm gives
          (TensorImage of size 4x3x224x224, TensorCategory([32, 20, 29, 35]))
        applying RandomResizedCropGPU gives
          (TensorImage of size 4x3x224x224, TensorCategory([32, 20, 29, 35]))
        applying LightingTfm gives
          (TensorImage of size 4x3x224x224, TensorCategory([32, 20, 29, 35]))
    CNN Learner
    cnn dls.device xla:1
    cnn dls.device MIDDLE
    cnn dls.device END
    FINE TUNE

      <progress value='0' class='' max='2' style='width:300px; height:20px; vertical-align: middle;'></progress>
      0.00% [0/2 00:00<00:00]

      epoch
      train_loss
      valid_loss
      accuracy
      time

      <progress value='17' class='' max='23' style='width:300px; height:20px; vertical-align: middle;'></progress>
      73.91% [17/23 04:18<01:31 3.8557]

                                              Learner#ALL_BATCHES
                                              Learner#ALL_BATCHES GEN enumerator!!!!
                                              Learner#ALL_BATCHES generate enumerator!!!!
     (07:15:25.013256)  DataLoader#DataLoader#DataLoader#__iter__                         0
     (07:15:25.014005)  DataLoader#DataLoader#DataLoader#__iter__ START FOR               1
     (07:15:25.097918)  DataLoader#DataLoader#DataLoader#__iter__ ----------------------- 1
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:15:37.173562)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:15:37.176881)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:15:39.316684)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:15:39.319828)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [0]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:15:54.172144)                   Learner#one_batch@BATCH 0
     (07:15:54.172765)                   Learner#one_batch@SPLIT 0
     (07:15:54.188308)                 Learner#one_batch@ 0 called begin_batch
     (07:16:02.418419)                 Learner#one_batch@ 0 called after_pred
     (07:16:02.419087)                 Learner#one_batch@ 0 did not return becasue yb.length == 0
     (07:16:02.429082)                 Learner#one_batch@ 0 called after_loss
     (07:16:02.429621)                 Learner#one_batch@ 0 did not return because we are training
     (07:16:11.455029)                 Learner#one_batch@ 0 called after_backward
     (07:16:27.284186)                 Learner#one_batch@ 0 called after_step
     (07:16:27.292841)                 Learner#one_batch@ 0 called zero_grad!!!!!!!
     (07:16:27.293264)                 0 WILL CALL after_batch
     (07:16:27.639723)                 0 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:16:27.642913)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:16:27.644895)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:16:27.645288)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:16:29.304716)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:16:29.305916)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [1]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:16:38.629329)                   Learner#one_batch@BATCH 1
     (07:16:38.629725)                   Learner#one_batch@SPLIT 1
     (07:16:38.638068)                 Learner#one_batch@ 1 called begin_batch
     (07:16:38.935730)                 Learner#one_batch@ 1 called after_pred
     (07:16:38.936368)                 Learner#one_batch@ 1 did not return becasue yb.length == 0
     (07:16:38.945960)                 Learner#one_batch@ 1 called after_loss
     (07:16:38.946394)                 Learner#one_batch@ 1 did not return because we are training
     (07:16:39.479736)                 Learner#one_batch@ 1 called after_backward
     (07:16:39.527306)                 Learner#one_batch@ 1 called after_step
     (07:16:39.535856)                 Learner#one_batch@ 1 called zero_grad!!!!!!!
     (07:16:39.536288)                 1 WILL CALL after_batch
     (07:16:39.836072)                 1 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:16:39.839167)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:16:39.860763)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:16:39.861283)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:16:41.507236)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:16:41.508168)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [2]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:16:50.692456)                   Learner#one_batch@BATCH 2
     (07:16:50.692858)                   Learner#one_batch@SPLIT 2
     (07:16:50.701309)                 Learner#one_batch@ 2 called begin_batch
     (07:16:50.992070)                 Learner#one_batch@ 2 called after_pred
     (07:16:50.992703)                 Learner#one_batch@ 2 did not return becasue yb.length == 0
     (07:16:51.002458)                 Learner#one_batch@ 2 called after_loss
     (07:16:51.002877)                 Learner#one_batch@ 2 did not return because we are training
     (07:16:51.503525)                 Learner#one_batch@ 2 called after_backward
     (07:16:51.551250)                 Learner#one_batch@ 2 called after_step
     (07:16:51.559538)                 Learner#one_batch@ 2 called zero_grad!!!!!!!
     (07:16:51.560124)                 2 WILL CALL after_batch
     (07:16:51.861375)                 2 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:16:51.864756)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:16:51.866781)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:16:51.867227)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:16:53.458107)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:16:53.458968)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [3]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:17:02.871614)                   Learner#one_batch@BATCH 3
     (07:17:02.872013)                   Learner#one_batch@SPLIT 3
     (07:17:02.880486)                 Learner#one_batch@ 3 called begin_batch
     (07:17:03.176678)                 Learner#one_batch@ 3 called after_pred
     (07:17:03.177407)                 Learner#one_batch@ 3 did not return becasue yb.length == 0
     (07:17:03.186884)                 Learner#one_batch@ 3 called after_loss
     (07:17:03.187308)                 Learner#one_batch@ 3 did not return because we are training
     (07:17:03.692511)                 Learner#one_batch@ 3 called after_backward
     (07:17:03.736704)                 Learner#one_batch@ 3 called after_step
     (07:17:03.744625)                 Learner#one_batch@ 3 called zero_grad!!!!!!!
     (07:17:03.745025)                 3 WILL CALL after_batch
     (07:17:04.045452)                 3 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:17:04.048594)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:17:04.062058)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:17:04.062474)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:17:05.695723)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:17:05.696649)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [4]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:17:15.173362)                   Learner#one_batch@BATCH 4
     (07:17:15.173872)                   Learner#one_batch@SPLIT 4
     (07:17:15.182551)                 Learner#one_batch@ 4 called begin_batch
     (07:17:15.475314)                 Learner#one_batch@ 4 called after_pred
     (07:17:15.475943)                 Learner#one_batch@ 4 did not return becasue yb.length == 0
     (07:17:15.485822)                 Learner#one_batch@ 4 called after_loss
     (07:17:15.486259)                 Learner#one_batch@ 4 did not return because we are training
     (07:17:15.994845)                 Learner#one_batch@ 4 called after_backward
     (07:17:16.060033)                 Learner#one_batch@ 4 called after_step
     (07:17:16.069190)                 Learner#one_batch@ 4 called zero_grad!!!!!!!
     (07:17:16.069716)                 4 WILL CALL after_batch
     (07:17:16.368443)                 4 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:17:16.371542)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:17:16.390493)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:17:16.390998)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:17:18.015838)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:17:18.017007)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [5]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:17:27.511329)                   Learner#one_batch@BATCH 5
     (07:17:27.511716)                   Learner#one_batch@SPLIT 5
     (07:17:27.519869)                 Learner#one_batch@ 5 called begin_batch
     (07:17:27.806226)                 Learner#one_batch@ 5 called after_pred
     (07:17:27.806851)                 Learner#one_batch@ 5 did not return becasue yb.length == 0
     (07:17:27.816301)                 Learner#one_batch@ 5 called after_loss
     (07:17:27.816708)                 Learner#one_batch@ 5 did not return because we are training
     (07:17:28.316518)                 Learner#one_batch@ 5 called after_backward
     (07:17:28.371950)                 Learner#one_batch@ 5 called after_step
     (07:17:28.379936)                 Learner#one_batch@ 5 called zero_grad!!!!!!!
     (07:17:28.380358)                 5 WILL CALL after_batch
     (07:17:28.679914)                 5 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:17:28.683002)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:17:28.699103)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:17:28.699491)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:17:30.329500)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:17:30.330374)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [6]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:17:39.871819)                   Learner#one_batch@BATCH 6
     (07:17:39.872232)                   Learner#one_batch@SPLIT 6
     (07:17:39.881148)                 Learner#one_batch@ 6 called begin_batch
     (07:17:40.168253)                 Learner#one_batch@ 6 called after_pred
     (07:17:40.168862)                 Learner#one_batch@ 6 did not return becasue yb.length == 0
     (07:17:40.178524)                 Learner#one_batch@ 6 called after_loss
     (07:17:40.178925)                 Learner#one_batch@ 6 did not return because we are training
     (07:17:40.679352)                 Learner#one_batch@ 6 called after_backward
     (07:17:40.741932)                 Learner#one_batch@ 6 called after_step
     (07:17:40.750410)                 Learner#one_batch@ 6 called zero_grad!!!!!!!
     (07:17:40.750818)                 6 WILL CALL after_batch
     (07:17:41.049912)                 6 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:17:41.053013)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:17:41.067639)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:17:41.068008)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:17:42.706829)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:17:42.707724)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [7]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:17:52.108578)                   Learner#one_batch@BATCH 7
     (07:17:52.108973)                   Learner#one_batch@SPLIT 7
     (07:17:52.117170)                 Learner#one_batch@ 7 called begin_batch
     (07:17:52.407579)                 Learner#one_batch@ 7 called after_pred
     (07:17:52.408273)                 Learner#one_batch@ 7 did not return becasue yb.length == 0
     (07:17:52.418342)                 Learner#one_batch@ 7 called after_loss
     (07:17:52.418863)                 Learner#one_batch@ 7 did not return because we are training
     (07:17:52.929617)                 Learner#one_batch@ 7 called after_backward
     (07:17:52.985071)                 Learner#one_batch@ 7 called after_step
     (07:17:52.993915)                 Learner#one_batch@ 7 called zero_grad!!!!!!!
     (07:17:52.994445)                 7 WILL CALL after_batch
     (07:17:53.294173)                 7 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:17:53.297280)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:17:53.312112)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:17:53.312535)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:17:54.952427)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:17:54.953320)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [8]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:18:04.588038)                   Learner#one_batch@BATCH 8
     (07:18:04.588458)                   Learner#one_batch@SPLIT 8
     (07:18:04.597155)                 Learner#one_batch@ 8 called begin_batch
     (07:18:04.880335)                 Learner#one_batch@ 8 called after_pred
     (07:18:04.880982)                 Learner#one_batch@ 8 did not return becasue yb.length == 0
     (07:18:04.890976)                 Learner#one_batch@ 8 called after_loss
     (07:18:04.891410)                 Learner#one_batch@ 8 did not return because we are training
     (07:18:05.389839)                 Learner#one_batch@ 8 called after_backward
     (07:18:05.444373)                 Learner#one_batch@ 8 called after_step
     (07:18:05.452203)                 Learner#one_batch@ 8 called zero_grad!!!!!!!
     (07:18:05.452610)                 8 WILL CALL after_batch
     (07:18:05.752579)                 8 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:18:05.755678)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:18:05.769239)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:18:05.769638)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:18:07.387748)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:18:07.389016)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [9]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:18:16.712972)                   Learner#one_batch@BATCH 9
     (07:18:16.713383)                   Learner#one_batch@SPLIT 9
     (07:18:16.721808)                 Learner#one_batch@ 9 called begin_batch
     (07:18:17.009228)                 Learner#one_batch@ 9 called after_pred
     (07:18:17.009853)                 Learner#one_batch@ 9 did not return becasue yb.length == 0
     (07:18:17.019329)                 Learner#one_batch@ 9 called after_loss
     (07:18:17.019736)                 Learner#one_batch@ 9 did not return because we are training
     (07:18:17.525829)                 Learner#one_batch@ 9 called after_backward
     (07:18:17.579974)                 Learner#one_batch@ 9 called after_step
     (07:18:17.588367)                 Learner#one_batch@ 9 called zero_grad!!!!!!!
     (07:18:17.588795)                 9 WILL CALL after_batch
     (07:18:17.889073)                 9 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:18:17.892220)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:18:17.911748)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:18:17.912169)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:18:19.543314)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:18:19.544552)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [10]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:18:28.829445)                   Learner#one_batch@BATCH 10
     (07:18:28.829846)                   Learner#one_batch@SPLIT 10
     (07:18:28.838416)                 Learner#one_batch@ 10 called begin_batch
     (07:18:29.122685)                 Learner#one_batch@ 10 called after_pred
     (07:18:29.123335)                 Learner#one_batch@ 10 did not return becasue yb.length == 0
     (07:18:29.133094)                 Learner#one_batch@ 10 called after_loss
     (07:18:29.133508)                 Learner#one_batch@ 10 did not return because we are training
     (07:18:29.643633)                 Learner#one_batch@ 10 called after_backward
     (07:18:29.701338)                 Learner#one_batch@ 10 called after_step
     (07:18:29.709617)                 Learner#one_batch@ 10 called zero_grad!!!!!!!
     (07:18:29.710074)                 10 WILL CALL after_batch
     (07:18:30.008983)                 10 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:18:30.012149)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:18:30.027987)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:18:30.028515)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:18:31.662151)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:18:31.663082)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [11]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:18:41.035114)                   Learner#one_batch@BATCH 11
     (07:18:41.035518)                   Learner#one_batch@SPLIT 11
     (07:18:41.044067)                 Learner#one_batch@ 11 called begin_batch
     (07:18:41.329021)                 Learner#one_batch@ 11 called after_pred
     (07:18:41.329682)                 Learner#one_batch@ 11 did not return becasue yb.length == 0
     (07:18:41.339736)                 Learner#one_batch@ 11 called after_loss
     (07:18:41.340241)                 Learner#one_batch@ 11 did not return because we are training
     (07:18:41.852346)                 Learner#one_batch@ 11 called after_backward
     (07:18:41.908890)                 Learner#one_batch@ 11 called after_step
     (07:18:41.917224)                 Learner#one_batch@ 11 called zero_grad!!!!!!!
     (07:18:41.917640)                 11 WILL CALL after_batch
     (07:18:42.218283)                 11 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:18:42.221410)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:18:42.237679)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:18:42.238223)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:18:43.871999)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:18:43.872796)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [12]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:18:53.458787)                   Learner#one_batch@BATCH 12
     (07:18:53.459203)                   Learner#one_batch@SPLIT 12
     (07:18:53.467751)                 Learner#one_batch@ 12 called begin_batch
     (07:18:53.759692)                 Learner#one_batch@ 12 called after_pred
     (07:18:53.760358)                 Learner#one_batch@ 12 did not return becasue yb.length == 0
     (07:18:53.770142)                 Learner#one_batch@ 12 called after_loss
     (07:18:53.770550)                 Learner#one_batch@ 12 did not return because we are training
     (07:18:54.281080)                 Learner#one_batch@ 12 called after_backward
     (07:18:54.336410)                 Learner#one_batch@ 12 called after_step
     (07:18:54.344573)                 Learner#one_batch@ 12 called zero_grad!!!!!!!
     (07:18:54.344983)                 12 WILL CALL after_batch
     (07:18:54.644160)                 12 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:18:54.647248)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:18:54.667883)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:18:54.668547)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:18:56.300111)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:18:56.301397)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [13]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:19:05.417989)                   Learner#one_batch@BATCH 13
     (07:19:05.418402)                   Learner#one_batch@SPLIT 13
     (07:19:05.426375)                 Learner#one_batch@ 13 called begin_batch
     (07:19:05.718329)                 Learner#one_batch@ 13 called after_pred
     (07:19:05.718960)                 Learner#one_batch@ 13 did not return becasue yb.length == 0
     (07:19:05.728667)                 Learner#one_batch@ 13 called after_loss
     (07:19:05.729084)                 Learner#one_batch@ 13 did not return because we are training
     (07:19:06.230183)                 Learner#one_batch@ 13 called after_backward
     (07:19:06.286636)                 Learner#one_batch@ 13 called after_step
     (07:19:06.294545)                 Learner#one_batch@ 13 called zero_grad!!!!!!!
     (07:19:06.294932)                 13 WILL CALL after_batch
     (07:19:06.594544)                 13 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:19:06.597763)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:19:06.613474)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:19:06.613864)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:19:08.192406)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:19:08.193634)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [14]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:19:17.507480)                   Learner#one_batch@BATCH 14
     (07:19:17.507878)                   Learner#one_batch@SPLIT 14
     (07:19:17.516244)                 Learner#one_batch@ 14 called begin_batch
     (07:19:17.807357)                 Learner#one_batch@ 14 called after_pred
     (07:19:17.807992)                 Learner#one_batch@ 14 did not return becasue yb.length == 0
     (07:19:17.817977)                 Learner#one_batch@ 14 called after_loss
     (07:19:17.818410)                 Learner#one_batch@ 14 did not return because we are training
     (07:19:18.335078)                 Learner#one_batch@ 14 called after_backward
     (07:19:18.393615)                 Learner#one_batch@ 14 called after_step
     (07:19:18.401612)                 Learner#one_batch@ 14 called zero_grad!!!!!!!
     (07:19:18.402011)                 14 WILL CALL after_batch
     (07:19:18.702800)                 14 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:19:18.705930)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:19:18.725490)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:19:18.726006)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:19:20.354525)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:19:20.355775)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [15]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:19:29.840756)                   Learner#one_batch@BATCH 15
     (07:19:29.841477)                   Learner#one_batch@SPLIT 15
     (07:19:29.849961)                 Learner#one_batch@ 15 called begin_batch
     (07:19:30.138806)                 Learner#one_batch@ 15 called after_pred
     (07:19:30.139467)                 Learner#one_batch@ 15 did not return becasue yb.length == 0
     (07:19:30.149060)                 Learner#one_batch@ 15 called after_loss
     (07:19:30.149466)                 Learner#one_batch@ 15 did not return because we are training
     (07:19:30.665098)                 Learner#one_batch@ 15 called after_backward
     (07:19:30.721612)                 Learner#one_batch@ 15 called after_step
     (07:19:30.729836)                 Learner#one_batch@ 15 called zero_grad!!!!!!!
     (07:19:30.730274)                 15 WILL CALL after_batch
     (07:19:31.030333)                 15 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:19:31.033464)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:19:31.048404)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:19:31.051620)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1
     (07:19:32.638757)  DataLoader#DataLoader#DataLoader#iterator to done!!!!
     (07:19:32.639667)  DataLoader#DataLoader#DataLoader#yielding                    3!!!! yield self.after_batch(xla:1 y xla:1) len of b is 2
    >>>> =======================================================================
    >>>> =======================================================================
    >>>>                                           Learner#ALL_BATCHES CALL ENTER [16]
    >>>> =======================================================================
    >>>> ======================================================================= OOOOOOOOOOOOOOOO call self.one_batch
     (07:19:42.138351)                   Learner#one_batch@BATCH 16
     (07:19:42.138759)                   Learner#one_batch@SPLIT 16
     (07:19:42.147651)                 Learner#one_batch@ 16 called begin_batch
     (07:19:42.452301)                 Learner#one_batch@ 16 called after_pred
     (07:19:42.452916)                 Learner#one_batch@ 16 did not return becasue yb.length == 0
     (07:19:42.462576)                 Learner#one_batch@ 16 called after_loss
     (07:19:42.462990)                 Learner#one_batch@ 16 did not return because we are training
     (07:19:42.972250)                 Learner#one_batch@ 16 called after_backward
     (07:19:43.040582)                 Learner#one_batch@ 16 called after_step
     (07:19:43.048877)                 Learner#one_batch@ 16 called zero_grad!!!!!!!
     (07:19:43.049373)                 16 WILL CALL after_batch
     (07:19:43.349970)                 16 called after_batch
    <<<< ======================================================================= OOOOOOOOOOOOOOOO
    <<<< =======================================================================
    <<<< =======================================================================
    <<<<                                           Learner#ALL_BATCHES CALL EXIT
    <<<< ====================================================================***
    <<<< ========================================*******************************
    <<<< ===========================********************************************
     (07:19:43.353072)  DataLoader#DataLoader#DataLoader#yielding                    4!!!!
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    ======================================================================================     XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     (07:19:43.369419)  DataLoader#DataLoader#DataLoader#__iter__ FOR FOR FOR FOR ---------------------------------------------
     (07:19:43.369791)  DataLoader#DataLoader#DataLoader#iterator to device from cpu y cpu to xla:1

```
print("done")
```

```

```