Here x is a tensor, f is an nn.Module model. Both are on cuda:
x = x.to(torch.device('cuda'))
f = f.to(torch.device('cuda'))

The main loop looks like this:
while(some condition):
      x = f(x)

I noticed that pretty quickly training slows down. I believe this feedforward mechanism could be a problem. Do I need to use a different approach, like create an intermediary tensor for x and input it into f?