computeEncoder Class — pytorch Architecture

Source Code

aten/src/ATen/native/mps/operations/MultiTensorApply.h lines 51–96

template <>
struct FusedSgdEncodingFunctor<true> {
  void operator()(id<MTLComputeCommandEncoder>& computeEncoder,
                  id<MTLBuffer>& tensorArgumentBuffer,
                  const MetadataArguments& metadata_arguments,
                  const double weight_decay,
                  const double momentum,
                  const double lr,
                  const double dampening,
                  const bool nesterov,
                  const bool maximize,
                  const bool is_first_step) const {
    mtl_setArgs(computeEncoder,
                tensorArgumentBuffer,
                metadata_arguments,
                weight_decay,
                momentum,
                lr,
                dampening,
                nesterov,
                maximize,
                is_first_step);
  }

  void operator()(id<MTLComputeCommandEncoder>& computeEncoder,
                  id<MTLBuffer>& tensorArgumentBuffer,
                  const MetadataArguments& metadata_arguments,
                  const double weight_decay,
                  const double momentum,
                  const at::Tensor& lr,
                  const double dampening,
                  const bool nesterov,
                  const bool maximize,
                  const bool is_first_step) const {
    mtl_setArgs(computeEncoder,
                tensorArgumentBuffer,
                metadata_arguments,
                weight_decay,
                momentum,
                lr,
                dampening,
                nesterov,
                maximize,
                is_first_step);
  }
};