linear.h

#pragma once

#include <torch/nn/cloneable.h>
#include <torch/nn/module.h>
#include <torch/nn/pimpl.h>
#include <torch/types.h>

#include <cstddef>
#include <vector>

namespace torch {
namespace nn {
struct TORCH_API LinearOptions {
  LinearOptions(int64_t in, int64_t out);
  TORCH_ARG(int64_t, in);
  TORCH_ARG(int64_t, out);
  TORCH_ARG(bool, with_bias) = true;
};

class TORCH_API LinearImpl : public Cloneable<LinearImpl> {
 public:
  LinearImpl(int64_t in, int64_t out) : LinearImpl(LinearOptions(in, out)) {}
  explicit LinearImpl(LinearOptions options);

  void reset() override;

  void pretty_print(std::ostream& stream) const override;

  Tensor forward(const Tensor& input);

  LinearOptions options;

  Tensor weight;

  Tensor bias;
};

TORCH_MODULE(Linear);

} // namespace nn
} // namespace torch