I am training a simple 2 layers MLP in an online learning setting where batch size and number of epoch are 1.
The input size is (5000000, 28) and the network is (28-100-2).
However, the training speed is very slow (13 seconds for 5000 instances) compare to my previous implementation in keras (5 seconds for 10000 instances).
Is there a way to improve the speed in online training ?
I put the code for custom dataset and training here, I did try training with GPU but the speed is significantly slower than CPU because batch size is 1 in this case.
class Stream(Dataset):
    def __init__(self, x,y):
        self.x = x
        self.y = y
        self.length = x.size(0)

    def __iter__(self):
        return self
    
    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        x_ = self.x[idx]
        y_ = self.y[idx]
        return x_, y_

stream = DataLoader(Stream(x_train, y_train), batch_size = 1, shuffle = False, num_workers = 8)
for j, (x,y) in enumerate(stream):
        loss_, acc_  = model.observe(x,y)
        loss += loss_
        acc += acc_