title: Select a channel in a 4D tensor without reducing dimensions using PyTorch
body: <p>I have 4D tensors for YUV image data using PyTorch -</p>
  <pre><code>a = torch.rand(N, C, H, W)
  </code></pre>
  <p>C here is 3 since there are YUV - 3 channels. I want to send each of the channels separately into a loss function which requires a 4D tensor to be inputted. How do I go about this?
  My approach -</p>
  <pre><code>y = a[:,0,:,:]
  u = a[:,1,:,:]
  v = a[:,2,:,:]
  </code></pre>
  <p>But by doing this, 4D tensors are being converted into 3D tensors causing an error during input to the loss function. How do I retain the dimensionality of the tensors?</p>
  <pre><code>y = torch.reshape(y, (N, 1, H, W))
  u = torch.reshape(u, (N, 1, H, W))
  v = torch.reshape(v, (N, 1, H, W))
  </code></pre>
  <p><code>torch.reshape</code> is there but it involves lot of hard-coding. Is there a better way to get this done?</p>