drbh

fix: refactors and adjust name

3683745 about 1 year ago

387 Bytes

#pragma once

#include <torch/torch.h>

namespace adam_atan2 {

void adam_atan2_cuda_impl_(
    std::vector<at::Tensor> params,
    std::vector<at::Tensor> grads,
    std::vector<at::Tensor> exp_avgs,
    std::vector<at::Tensor> exp_avg_sqs,
    std::vector<at::Tensor> state_steps,
    const double lr,
    const double beta1,
    const double beta2,
    const double weight_decay);

}

	#pragma once

	#include <torch/torch.h>

	namespace adam_atan2 {

	void adam_atan2_cuda_impl_(
	std::vector<at::Tensor> params,
	std::vector<at::Tensor> grads,
	std::vector<at::Tensor> exp_avgs,
	std::vector<at::Tensor> exp_avg_sqs,
	std::vector<at::Tensor> state_steps,
	const double lr,
	const double beta1,
	const double beta2,
	const double weight_decay);

	}