Hello All,
I am trying to use TPUs on a Kaggle Notebook and installing Pytorch XLA as follows -
!pip3 install mkl

!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py

!python3 pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev

However once Pytorch XLA is installed, It does not let me use numpy functions like numpy.uint8 or for that matter even torch.from_numpy() does not work. The error always says - Numpy not Available
The whole stack is as follows -
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-abfcbbc939b0> in <module>
   1026                                  segmentation_Maps='/kaggle/input/pascal-voc/VOC2012/SegmentationClass/') 
   1027 dataloader = DataLoader(dataset, batch_size=5)
-> 1028 for _, data in enumerate(dataloader):
   1029     i = data['image']
   1030     gt = data['ground_truth']

/opt/conda/lib/python3.7/site-packages/torch/utils/data/dataloader.py in __next__(self)
    519             if self._sampler_iter is None:
    520                 self._reset()
--> 521             data = self._next_data()
    522             self._num_yielded += 1
    523             if self._dataset_kind == _DatasetKind.Iterable and \

/opt/conda/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _next_data(self)
    559     def _next_data(self):
    560         index = self._next_index()  # may raise StopIteration
--> 561         data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    562         if self._pin_memory:
    563             data = _utils.pin_memory.pin_memory(data)

/opt/conda/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py in fetch(self, possibly_batched_index)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

/opt/conda/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py in <listcomp>(.0)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

<ipython-input-1-abfcbbc939b0> in __getitem__(self, item)
    939         print(mask.shape)
    940         image = Image.fromarray(np.uint8(image)).convert('RGB')
--> 941         mask = torch.from_numpy(np.uint8(mask))
    942 
    943         image = self.transforms(image)

RuntimeError: Numpy is not available

I have no clue what exactly is going on here.  Does anybody know how to solve it ?
PS - Pytorch XLA updates PyTorch to a nightly 1.9 version, so maybe that causes an issue  ?
TIA