Running the example in Google Colab on Mac M1 GPU

[wayneisaacuy]

Hi,

Has anyone tried running this code using a Mac M1 GPU? I am using torch version 2.0.1. I edited the code in this repo to enable using MPS instead of CUDA, i.e. I modified the_acquire_device method to:

def _acquire_device(self):
        if self.args.use_gpu:
            if torch.cuda.is_available():
                os.environ["CUDA_VISIBLE_DEVICES"] = str(
                    self.args.gpu) if not self.args.use_multi_gpu else self.args.devices
                device = torch.device('cuda:{}'.format(self.args.gpu))
                print('Use GPU: cuda:{}'.format(self.args.gpu))
            elif getattr(torch,'has_mps',False):
                os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 
                device = "mps"
                print('Use GPU: mps')
        else:
            device = torch.device('cpu')
            print('Use CPU')
        return device

I tested the ETTh1 example in the provided google notebook. I used the same parameters. However, my results under the 'Visualization' section are really off if I use the Mac M1 GPU via MPS. If I use CPU, my results look similar to the plots in the provided Google Colab notebook. Does anyone know if MPS has any bugs for the functions used in the network architecture for informer?

Thanks!