Repository: pytorch
Issue: libtorch: runtime error when iterating batch of dataloader
Labels: module: cpp,triaged

### 🐛 Describe the bug

Hello, I'm trying to implement train code with libtorch(pytorch c++ api). But the runtime error always raises when it goes to iterate batch of dataloader:
``train.cpp``
```cpp
auto train_data = UDGDataset(traindata, V, L, C).map(Stack<UDGExample<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> >());
		auto valid_data = UDGDataset(validdata, V, L, C).map(Stack<UDGExample<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> >());
		auto test_data = UDGDataset(testdata, V, L, C).map(Stack<UDGExample<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> >());
		auto trainloader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(std::move(train_data), batch_size);
		auto validloader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(valid_data), batch_size);
		auto testloader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(test_data), batch_size);
		auto net = GnnModelUndirected(dim, C, hidden_size, L);
		torch::optim::SGD* optimizer = new torch::optim::SGD(net->parameters(), lr = lr);
		torch::nn::CrossEntropyLoss criterion;
		float old_acc = 0.0;
		float old_vacc = 0.0;
		for (int i = 0; i < epoch; ++i) {
			int num = 0;
			int correct_num = 0;
			int total_err1 = 0;
			int total_err0 = 0;
			net->train();
			for (auto& batch : *trainloader) {
				optimizer->zero_grad();
				auto out = net->forward(emb, batch.one_hot_s, batch.one_hot_t, batch.one_hot_p);
				auto loss = criterion(out, batch.y);
				loss.backward();
				optimizer->step();
				auto res = torch::argmax(out, 1);
				int batchsize = res.size(0);
				num += batchsize;
				int cnt = 0;
				int err1 = 0;
				int err0 = 0;
				for (int j = 0; j < batchsize; ++j) {
					if (res[j].item<int>() == batch.y[j].item<int>()) {
						cnt++;
					}
					else if (res[j].item<int>() == 1) {
						err1++;
					}
					else {
						err0++;
					}
				}
				logging << "accuracy: " << std::fixed << std::setprecision(4) << (float)cnt / batchsize << std::endl;
				logging << "err1: " << err1 << std::endl;
				logging << "err0: " << err0 << std::endl;
				correct_num += cnt;
				total_err1 += err1;
				total_err0 += err0;
			}
         // ...et cetera...
```
``model.hpp``
```cpp
#pragma once
#include<torch/torch.h>

//using namespace torch::indexing;

struct GnnModelUndirectedImpl : torch::nn::Module {
	GnnModelUndirectedImpl(int input_dim, int label_num, int hidden_size, int path_length): 
		gru(torch::nn::GRU(torch::nn::GRUOptions(label_num, hidden_size).num_layers(1).batch_first(true))),
		linear(torch::nn::Linear(hidden_size, 2)),
		relu(torch::nn::ReLU()),
		softmax(torch::nn::Softmax(torch::nn::SoftmaxOptions(1))),
		linear_vec(torch::nn::Linear(input_dim* path_length, hidden_size))
	{
		//auto gru = register_module("gru", torch::nn::GRU(torch::nn::GRUOptions(label_num, hidden_size).num_layers(1).batch_first(true)));
		//auto linear = register_module("linear", torch::nn::Linear(hidden_size, 2));
		//auto relu = register_module("relu", torch::nn::ReLU());
		//auto softmax = register_module("softmax", torch::nn::Softmax(torch::nn::SoftmaxOptions(1)));
		//auto linear_vec = register_module("linear_vec", torch::nn::Linear(input_dim * path_length, hidden_size));
	}

	torch::Tensor forward(torch::Tensor vec, torch::Tensor one_hot_s, torch::Tensor one_hot_t, torch::Tensor one_hot_p) {
		auto vec1 = linear_vec(vec);
		vec1 = relu(vec1);
		auto s = torch::mm(one_hot_s, vec);
		auto t = torch::mm(one_hot_t, vec);
		auto p = relu(std::get<0>(gru(one_hot_p)));
		p = p.index({ torch::indexing::Slice(torch::indexing::None, -1, torch::indexing::None) });
		auto sp = s * p;
		auto tp = t * p;
		auto r = sp + tp;
		auto out = linear(r);
		out = softmax(out);
		return out;
	}

	torch::nn::GRU gru;
	torch::nn::Linear linear;
	torch::nn::ReLU relu;
	torch::nn::Softmax softmax;
	torch::nn::Linear linear_vec;
};

TORCH_MODULE(GnnModelUndirected);

template <typename S = torch::Tensor, typename T = torch::Tensor, typename P = torch::Tensor, typename Y = torch::Tensor> struct UDGExample {
	S one_hot_s;
	T one_hot_t;
	P one_hot_p;
	Y y;
	UDGExample() = default;
	UDGExample(S s, T t, P p, Y y):one_hot_s(std::move(s)), one_hot_t(std::move(t)), one_hot_p(std::move(p)), y(std::move(y)){}
};

template <typename ExampleType = UDGExample<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> >
struct Stack : public torch::data::transforms::Collation<ExampleType> {
	ExampleType apply_batch(std::vector<ExampleType> examples) override {
		std::vector<torch::Tensor> one_hot_s, one_hot_t, one_hot_p, y;
		one_hot_s.reserve(examples.size());
		one_hot_t.reserve(examples.size());
		one_hot_p.reserve(examples.size());
		y.reserve(examples.size());
		for (auto& example : examples) {
			one_hot_s.push_back(std::move(example.one_hot_s));
			one_hot_t.push_back(std::move(example.one_hot_t));
			one_hot_p.push_back(std::move(example.one_hot_p));
			y.push_back(std::move(example.y));
		}
		return { torch::stack(one_hot_s), torch::stack(one_hot_t), torch::stack(one_hot_p), torch::stack(y)};
	}
};

class UDGDataset : public torch::data::Dataset<UDGDataset, UDGExample<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> > {
private:
	std::vector<std::vector<int> > Data;
	int N;
	int V;
	int L;
	int C;
public:
	UDGDataset(std::vector<std::vector<int> >& data, int v, int l, int c) {
		Data = data;
		N = data.size();
		V = v;
		L = l;
		C = c;
	}
	UDGExample<> get(size_t i) override {
		torch::Tensor one_hot_s = torch::zeros({ V });
		torch::Tensor one_hot_t = torch::zeros({ V });
		torch::Tensor one_hot_p = torch::zeros({ L, C });
		torch::Tensor y = torch::zeros({ 1 }, torch::dtype(torch::kLong));
		int s = Data[i][0];
		int t = Data[i][1];
		int l = Data[i][2];
		one_hot_s[s] = 1;
		one_hot_t[t] = 1;
		for (int j = 0; j < l; ++j) {
			one_hot_p[j][Data[i][3 + j]] = 1;
		}
		y[0] = Data[i][Data[i].size() - 1];
		return { one_hot_s.clone(), one_hot_t.clone(), one_hot_p.clone(), y.clone() };
	}
	torch::optional<size_t> size() const override {
		return N;
	}
};
```

The error raises in `.\libtorch\include\torch\csrc\api\include\torch\data\dataloader\stateless.h`, on `line 76`, the assert failed:
![image](https://github.com/pytorch/pytorch/assets/46044770/e8e5b11b-b5fa-4af8-92fe-dfbc5298b1e5)

The error information says that is a c10::error.

Should any further information be needed, please tell me. 

Hope for solution. Thank you.


### Versions

from pytorch download page, choose preview(Nightly), Windows, LibTorch, C++/Java, CPU, downloaded from `https://download.pytorch.org/libtorch/nightly/cpu/libtorch-win-shared-with-deps-latest.zip`

cc @jbschlosser