text: #pragma once
  
  #include <torch/data/transforms/base.h>
  
  #include <functional>
  #include <utility>
  #include <vector>
  
  namespace torch {
  namespace data {
  namespace transforms {
  
  /// A `BatchTransform` that applies a user-provided functor to a batch.
  template <typename Input, typename Output = Input>
  class BatchLambda : public BatchTransform<Input, Output> {
   public:
    using typename BatchTransform<Input, Output>::InputBatchType;
    using typename BatchTransform<Input, Output>::OutputBatchType;
    using FunctionType = std::function<OutputBatchType(InputBatchType)>;
  
    /// Constructs the `BatchLambda` from the given `function` object.
    explicit BatchLambda(FunctionType function)
        : function_(std::move(function)) {}
  
    /// Applies the user-provided function object to the `input_batch`.
    OutputBatchType apply_batch(InputBatchType input_batch) override {
      return function_(std::move(input_batch));
    }
  
   private:
    FunctionType function_;
  };
  
  // A `Transform` that applies a user-provided functor to individual examples.
  template <typename Input, typename Output = Input>
  class Lambda : public Transform<Input, Output> {
   public:
    using typename Transform<Input, Output>::InputType;
    using typename Transform<Input, Output>::OutputType;
    using FunctionType = std::function<Output(Input)>;
  
    /// Constructs the `Lambda` from the given `function` object.
    explicit Lambda(FunctionType function) : function_(std::move(function)) {}
  
    /// Applies the user-provided function object to the `input`.
    OutputType apply(InputType input) override {
      return function_(std::move(input));
    }
  
   private:
    FunctionType function_;
  };
  
  } // namespace transforms
  } // namespace data
  } // namespace torch
  