It seems that TensorDataset will squeeze the dimension of the tensor_data.
I tried to write a dummy code
tensor_data = torch.randn(10,).unsqueeze(0).unsqueeze(0)
custom_dataset = torch.utils.data.TensorDataset(tensor_data)
print(custom_dataset[0][0].shape) 
>> [1, 10]