Program Listing for File mnist.h¶
↰ Return to documentation for file (torch/csrc/api/include/torch/data/datasets/mnist.h
)
#pragma once
#include <torch/data/datasets/base.h>
#include <torch/data/example.h>
#include <torch/types.h>
#include <torch/csrc/Export.h>
#include <cstddef>
#include <string>
namespace torch {
namespace data {
namespace datasets {
class TORCH_API MNIST : public Dataset<MNIST> {
public:
enum class Mode { kTrain, kTest };
explicit MNIST(const std::string& root, Mode mode = Mode::kTrain);
Example<> get(size_t index) override;
optional<size_t> size() const override;
// NOLINTNEXTLINE(bugprone-exception-escape)
bool is_train() const noexcept;
const Tensor& images() const;
const Tensor& targets() const;
private:
Tensor images_, targets_;
};
} // namespace datasets
} // namespace data
} // namespace torch