Combining Tensor Matrix and Sparse Matrix for Splitting Data in PyTorch: A Custom Dataset Approach