I have an input tensor of size (batch_size, X, Y) and need to pass it though the forward step of my custom model.
At high level in the forward step:
I loop over each batch and send the inner tensor of shape (X, Y) to another model that gives me something of shape (X,Z).
Then I need to do the average over X and assign this result to each batch to get a final tensor of shape (batch_size, Z).
I can do it with for-loop, but I think it might be inefficient to run on a GPU (also, I am getting some CUDA error right now). A naive-example to better explain the problem.
def forward(self, in_tensor) """ in_tensor something like in_tensor = tensor([[[ 1., 2., 3.], [ 4., 5., 6.]], [[ 1., 2., 3.], [ 4., 5., 6.]]], dtype=torch.float) # shape [2,2,3] """ res =  for batch, el in enumerate(in_tensor): # el has shape [2,3] tmp_res = other_model(el) # other model return something with 10 features for each input # tmp res has shape (2, 10) res.append(torch.mean(tmp_res, dim=0)) res = torch.stack(res, dim=0) ...
I am wondering if there is a way to avoid the for-loop.