I don’t completely understand the issue and the posted shapes and I get the expected output as [batch_size, 1024, 2] and [batch_size, 1024, 1] as seen here:
class CustomDataset(Dataset):
    def __init__(self, x_tensor, y_tensor):
        self.xdomain = x_tensor
        self.ydomain = y_tensor
        
    def __getitem__(self, index):
        return (self.xdomain[index], self.ydomain[index])
    
    def __len__(self):
        return len(self.xdomain)
    
    
x_tensor = torch.randn(1024, 1024, 2)
y_tensor = torch.randn(1024, 1024, 1)
dataset = CustomDataset(x_tensor, y_tensor)
loader = DataLoader(dataset, batch_size=2)

for x, y in loader:
    print(x.shape, y.shape)
> torch.Size([2, 1024, 2]) torch.Size([2, 1024, 1])
  torch.Size([2, 1024, 2]) torch.Size([2, 1024, 1])
  torch.Size([2, 1024, 2]) torch.Size([2, 1024, 1])
  [...]