Thank you for checking it out!
Just run your code as is (w/o any modifications, so it’s cuda tensor) and:
on sizes 1:
batch_size = 300
height = 200 
width = 200 

got
0.028152191638946535
0.0026354169845581054

on sizes 2:
batch_size = 8 
height = 1024 
width = 1024 

got
0.0010659313201904297
0.05072437047958374

and on sizes 3:
batch_size = 2
height = 4096 
width = 4096 

got
0.0024710512161254883
0.7927418446540833

So, it gets slower when h and w is increasing, isn’t it weird?
Update 1:
same situation on torch version 1.5.0