@sbelharbi
Skimmed through your question. Seems to me that everything is working properly, and the problem is that you don’t (yet) have a correct mental model of what’s going on under the hood.
Let’s start with the bad news - I am quite certain that your forward pass takes much more time than you think. Probably as much time as you think the ‘if’ statement consumes. This time is just being (correctly) hidden from you by the CUDA system. What you have measured as your forward pass timing is just the time required to queue the work onto the CUDA driver and return to your CPU-based program execution.
The good news is that nothing weird is happening and everything you are trying to do is perfectly achievable. With extra work, you can also implement fine grained synchronization avoiding brute force torch.cuda.synchronize(), which is the concurrency equivalent of using a nuke in the battlefield.
Of course you will have to deal with the real run-time of the computation.
Unfortunately the Torch documentation is extremely sparse in regard to synchronization issues and mostly assumes that users are either:

naive, sequential CPU-dwellers who transfer all data to the GPU and then do all computation on the GPU from there on
experienced, scarred CUDA programmers who know everything there is to know about asynchronous kernel execution.

Your case seems to fall in between these categories so you have been seeing what seems like concurrency voodoo.
The issue of synchronization and async execution is much too big for a post here so I will just give a few basic principles, hopefully these and what I’ve written above will give enough information to seek a solution:


Every time you execute a torch operation on the GPU, it is “non-blocking”; in the sense that GPU is told to do some work in the future, and execution resumes on the Python CPU thread almost immediately. This frees up your Python program to keep generating future GPU work, or to prepare new data while the GPU is doing its thing.


The reason this asynchronous operation doesn’t cause immediate mayhem is that work is queued on something called a CUDA Stream; on the stream everything proceeds in the order you have queued it. The reason most torch users are not aware of this is that there is a “default stream”, and unless the user requests something else, all GPU operations are queued on it, making everything look sequential and calm, as long as your tensors remain in CUDA land.


All of this Stream business is a CUDA-only thing. Therefore it applies only to stuff happening on the ‘cuda’ device. On the CPU there are is no such thing as streams and everything is simply synchronous with your Python program, just as most users expect.


So we have the GPU kingdom where everything is sequential because it’s on the default stream, and the CPU realm where everything is sequential because it is Python program synchronous. Each of these two domains maintains its own order, and they are largely independent of each other. This of course leaves the border between them as the problem area.


By default when you send a tensor across the boundary (nonblocking=False), the operation is synchronous on both sides, i.e. the copy is queued on the default CUDA stream AND your Python program is blocked until the tensor is completely moved to the other domain. This is nice and safe. And wastes precious time. This is why torch came up with nonblocking=True.


When you use nonblocking transfers, the copy work is queued on the default CUDA stream but the CPU is free to proceed. If you are doing cpu->GPU transfers, then most likely what happens next is that your Python program will queue further GPU operations for the tensor you have just uploaded. This is ok because the compute work is queued on the same stream with the copy work, and placed after it, so everything will end up blissfully fine.


However, when you move data from the GPU to the CPU with nonblocking=True, there is nothing there to protect you implicitly. It will queue the copy on the CUDA default stream, but you Python program is free to proceed. Quite likely your GPU is still busy doing computations that you have previously queued, and hasn’t even “heard” of the CPU-side tensor you are expecting to do further compute on. Therefore that tensor is going to contain zeros or garbage when your CPU starts crunching it.


To properly use a GPU result on the CPU side with a non-blocking transfer, you will need to make sure that the copy to the CPU has completed before you start consuming the data.
Before discussing how this can be accomplished, you have to ask yourself whether this is at all useful to you. In many (simple) cases there is nothing useful for the CPU to do until the compute result from the GPU is available. If this is the case, the easiest way to ensure the copy is complete is to do a blocking transfer (set nonblocking=False). Since the copy work is queued on the default CUDA stream, the copy will begin after all computation on the data has finished. Since Python is blocked for it, you CPU computation will only proceed after the copy has finished. This is also when you need to stop the timer if you want to measure how much time the GPU computation really took.
If you decide that you need more advanced asynchronicity in you program, I advise you to study some more about CUDA kernel launching. Keep in mind that every torch operation on Tensors on the GPU is equivalent to a sequence of one or more CUDA kernel launches in succession. You will need to be familiar with the following CUDA concepts and their torch interfaces: Stream, default stream, Event, overlapped computation and data transfer.
I hope this helps.