title: How to move PyTorch model to GPU on Apple M1 chips?
body: <p>On 18th May 2022, PyTorch <a href="https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/" rel="nofollow noreferrer">announced support</a> for GPU-accelerated PyTorch training on Mac.</p>
  <p>I followed the following process to set up PyTorch on my Macbook Air M1 (using miniconda).</p>
  <pre><code>conda create -n torch-nightly python=3.8 
  
  $ conda activate torch-nightly
  
  $ pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
  </code></pre>
  <p>I am trying to execute a script from Udacity's Deep Learning Course available <a href="https://github.com/udacity/deep-learning-v2-pytorch/blob/master/dcgan-svhn/DCGAN_Solution.ipynb" rel="nofollow noreferrer">here</a>.</p>
  <p>The script moves the models to GPU using the following code:</p>
  <pre class="lang-py prettyprint-override"><code>G.cuda()
  D.cuda()
  </code></pre>
  <p>However, this will not work on M1 chips, since there is no CUDA.</p>
  <p>If we want to move models to M1 GPU and our tensors to M1 GPU, and train entirely on M1 GPU, what should we be doing?</p>
  <hr />
  <p>If Relevant: <code>G</code> and <code>D</code> are Discriminator and Generators for GAN's.</p>
  <pre class="lang-py prettyprint-override"><code>class Discriminator(nn.Module):
  
      def __init__(self, conv_dim=32):
          super(Discriminator, self).__init__()
          self.conv_dim = conv_dim
          # complete init function
          self.cv1 = conv(in_channels=3, out_channels=conv_dim, kernel_size=4, stride=2, padding=1, batch_norm=False)            # 32*32*3  -&gt; 16*16*32
          self.cv2 = conv(in_channels=conv_dim, out_channels=conv_dim*2, kernel_size=4, stride=2, padding=1, batch_norm=True)    # 16*16*32 -&gt; 8*8*64
          self.cv3 = conv(in_channels=conv_dim*2, out_channels=conv_dim*4, kernel_size=4, stride=2, padding=1, batch_norm=True)  # 8*8*64   -&gt; 4*4*128
          self.fc1 = nn.Linear(in_features = 4*4*conv_dim*4, out_features = 1, bias=True)
          
  
      def forward(self, x):
          # complete forward function
          out = F.leaky_relu(self.cv1(x), 0.2)
          out = F.leaky_relu(self.cv2(x), 0.2)
          out = F.leaky_relu(self.cv3(x), 0.2)
          out = out.view(-1, 4*4*conv_dim*4)
          out = self.fc1(out)
          return out    
  
  D = Discriminator(conv_dim)
  
  class Generator(nn.Module):    
      def __init__(self, z_size, conv_dim=32):
          super(Generator, self).__init__()
          self.conv_dim = conv_dim
          self.z_size = z_size
          # complete init function
          self.fc1 = nn.Linear(in_features = z_size, out_features = 4*4*conv_dim*4)
          self.dc1 = deconv(in_channels = conv_dim*4, out_channels = conv_dim*2, kernel_size=4, stride=2, padding=1, batch_norm=True)
          self.dc2 = deconv(in_channels = conv_dim*2, out_channels = conv_dim, kernel_size=4, stride=2, padding=1, batch_norm=True)
          self.dc3 = deconv(in_channels = conv_dim, out_channels = 3, kernel_size=4, stride=2, padding=1, batch_norm=False)
  
      def forward(self, x):
          # complete forward function
          x = self.fc1(x)
          x = x.view(-1, conv_dim*4, 4, 4)
          x = F.relu(self.dc1(x))
          x = F.relu(self.dc2(x))
          x = F.tanh(self.dc3(x))
          return x
  
  G = Generator(z_size=z_size, conv_dim=conv_dim)
  </code></pre>