I tried following code, I simplified it a bit, seems to work fine, can you try ?
class ClassificationDataset : public torch::data::Dataset<ClassificationDataset>
{
private:
	std::vector<std::tuple<std::string, int>> m_data;

	torch::Tensor LoadImageToTensor(const std::string& path)
	{
		cv::Mat img = cv::imread(path);
		cv::resize(img, img, cv::Size(224, 224));

		torch::Tensor tensor_image = torch::from_blob(img.data, { img.rows, img.cols, 3 }, at::kByte);
		tensor_image = tensor_image.toType(at::kFloat);
		tensor_image = tensor_image.div_(255);
		tensor_image = tensor_image.permute({ 2, 0, 1 });

		return tensor_image;
	}

public:
	ClassificationDataset(const std::string& path) 
	{
		for (const auto& dir : std::filesystem::directory_iterator(path))
		{
			int label = std::atoi(dir.path().stem().string().c_str());

			for (const auto& file : std::filesystem::directory_iterator(dir))
			{
				m_data.push_back({ file.path().string(), label });
			}
		}
	};

	torch::data::Example<torch::Tensor, torch::Tensor> get(size_t index) override 
	{
		auto [path, label] = m_data[index];

		torch::Tensor sample_img = LoadImageToTensor(path);
		torch::Tensor sample_label = torch::tensor(label);

		return { sample_img, sample_label };
	};

	torch::optional<size_t> size() const override 
	{
		return m_data.size();
	};
};

void main(int argc, char** argv)
{
	try
	{
		auto dataset_train = ClassificationDataset("D:\\train").map(torch::data::transforms::Stack<>());
		auto data_loader_train = torch::data::make_data_loader(std::move(dataset_train), torch::data::DataLoaderOptions().batch_size(5).workers(4));

		auto dataset_val = ClassificationDataset("D:\\val").map(torch::data::transforms::Stack<>());
		auto data_loader_val = torch::data::make_data_loader(std::move(dataset_val), torch::data::DataLoaderOptions().batch_size(2).workers(4));

		for (auto& batch : *data_loader_train)
		{
			auto data = batch.data;
			auto labels = batch.target;

			std::cout << data.sizes() << std::endl;
			std::cout << labels.sizes() << std::endl;
		}

		for (auto& batch : *data_loader_val)
		{
			auto data = batch.data;
			auto labels = batch.target;

			std::cout << data.sizes() << std::endl;
			std::cout << labels.sizes() << std::endl;
		}
	}
	catch (std::runtime_error& e)
	{
		std::cout << e.what() << std::endl;
	}
	catch (const c10::Error& e)
	{
		std::cout << e.msg() << std::endl;
	}

	system("PAUSE");
}