I have two machines that I need to check my code across one is Ubuntu 18.04 and the other is Ubuntu 20.04.
In Ubuntu 18.04 I have:
$ python
Python 3.8.0 (default, Dec  9 2021, 17:53:27)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'1.9.1+cu102'

However, the code that works in Ubuntu 20.04, throws this error:
    with torch.autocast('cuda'):
AttributeError: module 'torch' has no attribute 'autocast'

I have this version of PyTorch on Ubuntu 20.04:
 python
Python 3.8.10 (default, Nov 26 2021, 20:14:08)
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'1.10.1+cu113'

Here’s the GPU info on Ubuntu 20.04:

Screen Shot 2022-02-01 at 2.45.51 AM814×566 59.7 KB

and here’s the GPU info on Ubuntu 18.04:

Screen Shot 2022-02-01 at 2.46.59 AM822×551 59.4 KB

Ideally I want the same code to run across two machines.