Torch: partition tensor

Question

I want to partition my dataset (10,000 50x50 RGB images) into two datasets. Something like:

X = torch.rand(10000, 3, 50, 50)
inds = torch.randperm(X:size(1))[{ { 1, nTrain } }]:long()
X_selected = X:index(1, inds)
X_remaining = X:delete(1, inds)

No matter what I google search, I just get Torch's GitHub documentation. How can I do this?

Answer 1

You can try this way

X = torch.rand(10000, 3, 50, 50)
inds = torch.randperm(X:size(1)):long()
train_inds = inds:narrow(1, 1, nTrain)
valid_inds = inds:narrow(1, nTrain + 1, X:size(1) - nTrain)
X_train = X:index(1, train_inds)
X_valid = X:index(1, valid_inds)