PyTorch not available error on SageMaker GPU docker though it is installed 
Hey @sivakhno,

how does your `requirements.txt` look like to install the `datasets` library and which version of it are you running? Can you try to install `datasets>=1.4.0`
I get en error when running data loading using SageMaker SDK

```
  File "main.py", line 34, in <module>
    run_training()
  File "main.py", line 25, in run_training
    dm.setup('fit')
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/core/datamodule.py", line 92, in wrapped_fn
    return fn(*args, **kwargs)
  File "/opt/ml/code/data_module.py", line 103, in setup
    self.dataset[split].set_format(type="torch", columns=self.columns)
  File "/opt/conda/lib/python3.6/site-packages/datasets/fingerprint.py", line 337, in wrapper
    out = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 995, in set_format
    _ = get_formatter(type, **format_kwargs)
File "/opt/conda/lib/python3.6/site-packages/datasets/formatting/__init__.py", line 114, in get_formatter
    raise _FORMAT_TYPES_ALIASES_UNAVAILABLE[format_type]
ValueError: PyTorch needs to be installed to be able to return PyTorch tensors.
```

when trying to execute dataset loading using this notebook  https://github.com/PyTorchLightning/pytorch-lightning/blob/master/notebooks/04-transformers-text-classification.ipynb, specifically lines 

```
self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns]
self.dataset[split].set_format(type="torch", columns=self.columns)
```

The SageMaker docker image used is 763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training:1.4.0-gpu-py3 .

By running container interactively I have checked that torch loading completes successfully by executing `https://github.com/huggingface/datasets/blob/master/src/datasets/config.py#L39`. 

Also as a first line in the data loading module I have 

```
import os
os.environ["USE_TF"] = "0" 
os.environ["USE_TORCH"] = "1" 
````

But unfortunately the error stills persists. Any suggestions would be appreciated as I am stack.
Many Thanks! 


PyTorch not available error on SageMaker GPU docker though it is installed  
 I get en error when running data loading using SageMaker SDK

```
  File "main.py", line 34, in <module>
    run_training()
  File "main.py", line 25, in run_training
    dm.setup('fit')
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/core/datamodule.py", line 92, in wrapped_fn
    return fn(*args, **kwargs)
  File "/opt/ml/code/data_module.py", line 103, in setup
    self.dataset[split].set_format(type="torch", columns=self.columns)
  File "/opt/conda/lib/python3.6/site-packages/datasets/fingerprint.py", line 337, in wrapper
    out = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 995, in set_format
    _ = get_formatter(type, **format_kwargs)
File "/opt/conda/lib/python3.6/site-packages/datasets/formatting/__init__.py", line 114, in get_formatter
    raise _FORMAT_TYPES_ALIASES_UNAVAILABLE[format_type]
ValueError: PyTorch needs to be installed to be able to return PyTorch tensors.
```

when trying to execute dataset loading using this notebook  https://github.com/PyTorchLightning/pytorch-lightning/blob/master/notebooks/04-transformers-text-classification.ipynb, specifically lines 

```
self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns]
self.dataset[split].set_format(type="torch", columns=self.columns)
```

The SageMaker docker image used is 763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-training:1.4.0-gpu-py3 .

By running container interactively I have checked that torch loading completes successfully by executing `https://github.com/huggingface/datasets/blob/master/src/datasets/config.py#L39`. 

Also as a first line in the data loading module I have 

```
import os
os.environ["USE_TF"] = "0" 
os.environ["USE_TORCH"] = "1" 
````

But unfortunately the error stills persists. Any suggestions would be appreciated as I am stack.
Many Thanks! 

 
 Hey @sivakhno,

how does your `requirements.txt` look like to install the `datasets` library and which version of it are you running? Can you try to install `datasets>=1.4.0`