I’m running into a weird performance issue with some blocks after a long kernel conv1d.  In the example below, the simple operation of picking random index from a batch is 1.44s with a long kernel.  The conv1d kernel performance is fine.  It’s only the subsequent block that is problematic.  This is experiment was done with v100 GPU on Ubuntu 18.04, cuda 11.1.  What’s the recommended fix?
import torch
from time import time

print('Torch version {}'.format(torch.__version__))


def test_conv_timing(K):
    print('Kernel: {}'.format(5))
    x = torch.randn(1, 512, 160000, dtype=torch.float32, device='cuda')
    kernel = torch.randn(512, 1, K, dtype=torch.float32, device='cuda')
    for i in range(10):
        dummy = torch.nn.functional.conv1d(x, kernel, groups=512)
        start_time = time()
        idx = torch.randint(0, 512, (1,))
        y = x[:, idx, :]
        print('{} elapsed time {}s'.format(i, time() - start_time))


test_conv_timing(K=5)
test_conv_timing(K=20000)

Output:
Kernel: 5
0 elapsed time 0.018301963806152344s
1 elapsed time 0.0017490386962890625s
2 elapsed time 0.0017266273498535156s
3 elapsed time 0.0017251968383789062s
4 elapsed time 0.001725912094116211s
5 elapsed time 0.001722574234008789s
6 elapsed time 0.0017278194427490234s
7 elapsed time 0.0017235279083251953s
8 elapsed time 0.0017247200012207031s
9 elapsed time 0.0017256736755371094s
Kernel: 20000
0 elapsed time 1.4438841342926025s
1 elapsed time 1.4291105270385742s
2 elapsed time 1.43113374710083s
3 elapsed time 1.4308841228485107s
4 elapsed time 1.4307475090026855s
5 elapsed time 1.430490255355835s
6 elapsed time 1.4306378364562988s
7 elapsed time 1.430649995803833s
8 elapsed time 1.4352474212646484s
9 elapsed time 1.4375996589660645s