Solved by ptrblck in post #2 
              
              
                You can use the context manager to enable/disable benchmark for a region: 
with torch.backends.cudnn.flags(enabled, benchmark, deterministic, allow_tf32):
    some_layer