Olliejp:

train_X = train_X.to('cuda')
train_Y = train_Y.to('cuda')
valid_X = valid_X.to('cuda')
valid_Y = valid_Y.to('cuda')



The main reason as I can see is you are trying to put the whole train and valid datasets to cuda device (GPU) at once. If your dataset is big enough it will simply not fit into GPU memory.
Second point is here:



 Olliejp:

train_loss = loss_fn(model(train_X.float()), train_Y.float())
val_loss = loss_fn(model(valid_X.float()), valid_Y.float())



You are trying to do a forward and backward pass of your model through whole datasets again.
My suggestion would be to work through official pytorch tutorials 10. And keep in mind we use pytorch dataloader to work on batches and not on the full dataset, mostly because we always have GPU memory constraints  and most of our datasets don’t fit into single GPU memory at once.