title: Applying a 2D function to a 4D tensor PyTorch
body: <p>I got a 2D function that takes a matrix - 2D tensor with shape (28, 28)
  and I got a tensor, lets say (64, 10, 28, 28) - it's a tensor that contains a batch of 64 images that passed through a (10 kernels) conv2d layer.</p>
  <p>Now, I want to activate on the last two dimentions of the tensor, the (28,28) bit, a 2D function.</p>
  <p>Now I did that in a very inefficient way:</p>
  <pre><code>def activation_func(input):
  
      for batch_idx in range(input.shape[0]):
          for channel_inx in range(input.shape[1]):       
              input[batch_idx][channel_inx] = 2D_function(input[batch_idx][channel_inx])
  
      return input
  </code></pre>
  <p>which is highly inefficient as I noticed.
  is there any way of doing this efficiently?</p>
  <p>I can write the entire code If necessary</p>
  <p>EDIT:</p>
  <pre><code>def 2D_function(input):
      global indices # yes I know, I will remove this global stuff later
      # indices = [(i, j) for i in range(1, 28, 4) for j in range(1, 28, 4)]
  
      for x, y in indices:
          relu_decision = relu(input[x, y]) # standard relu - relu(x)=(x&gt;1)*x
          if not relu_decision:
              # zero out the patch
              input[x - 1: x + 3, y - 1: y + 3] = 0
      return input
  
  </code></pre>