I found some solution for this problem in tutorial. But, I’m not sure. I want to check my code is working right whether or not. This example code is conceptual code.

kernel.cu

#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/core/TensorAccessor.h>

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
using at::Tensor;

Tensor CustomLossFunction(Tensor &grad, const Tensor &A, const long long N){
    auto B = at::zeros(N).to(at::kCUDA);
    LossFunction<<<nBlock, nThread>>>(grad.data<float>(), A.data<float>(), B.data<float>(), N);
   return B;
}
template<typename T>
__global__ void LossFunction(T *grad, T *B, const T *A, const long long N){
    const size_t tid = blockDim.x * gridDim.x + blockIdx.x * blockDim.x + threadIdx.x;
    if(tid < N){
         C[tid] = A[tid] * A[tid] * 2.f;
         grad[tid] = A[tid] * 4.f;
    }
}


main.cpp

#include <ATen/ATen.h>
#include <torch/torch.h>
#include "main.h"
using at::Tensor;

Tensor CustomLossFunction(Tensor &grad, const Tensor &A, const long long N);
int main(){
   Net model;
   model.to(torch::kCUDA);
   torch::optim::SGD optimizer(model.parameters(), torch::optim::SGDOptions(0.001).momentum(0.5));
   auto model_in = DataLoader(fileName).to(torch::kCUDA);
   auto Exact = DataLoader(fileName).to(torch::kCUDA);
   
   optimizer.zero_grad();
   auto grad = torch::zeros(10).to(Torch::kCUDA);
   auto Y = model.forward(model_in);
   auto Loss = CustomLossFunction(grad, Y, 10).mean();
   Y.backward(grad / 10);
   optimizer.step();
}

Updating parameters using custom loss function.