How can I make torch.cat faster?
Below is a simplified version of what I want to do:
import torch
import time

# Create dummy tensors and save them in my_list
my_list = [[]] * 100
for i in range(len(my_list)):
    my_list[i] = torch.randint(0, 1000000000, (100000, 256))
concat_list = torch.tensor([])

# I want to concat two consecutive tensors in my_list
tic = time.time()
for i in range(0, len(my_list), 2):
    concat_list = torch.cat((concat_list, my_list[i]))
    concat_list = torch.cat((concat_list, my_list[i+1]))
    # Do some work at CPU with concat_list
    concat_list = torch.tensor([]) # Empty concat_list
print('time: ', time.time() - tic) # It takes 3.5 seconds in my environment

Is there any way to make above tensor concatenation faster?
I tried to send my_list[i], my_list[i+1], and concat_list to GPU and do the torch.cat function in the device, but I then have to send concat_list back to CPU to do &quot;some work&quot; that I've written above. This takes more time due to frequent GPU-CPU data transfer.
I've also tested converting tensors to lists to do the concatenation with basic Python lists, but this approach was way slower than a simple torch.cat approach.
I've heard that using DataLoader with customized collate_fn can enable concatenation, but I don't know how to implement it.
Is there any faster method possible?

Your code takes around 11 seconds on my PC. The following takes 4.1 seconds:
# Create dummy tensors and save them in my_list
my_list = [[]] * 100
for i in range(len(my_list)):
    my_list[i] = torch.randint(0, 1000000000, (100000, 256))
tic = time.time()
my_list = torch.stack(my_list)

# I want to concat two consecutive tensors in my_list
for i in range(0, len(my_list), 2):
    concat_list = my_list[i:i+2]
    # Do some work at CPU with concat_list
print('time: ', time.time() - tic) # It takes 3.5 seconds in my environment

