Hi!
The CUDA version is implemented here: https://github.com/pytorch/pytorch/blob/5d7770948587bb5f19b929200f20e2eaf7074e5c/aten/src/THCUNN/generic/DistKLDivCriterion.cu 35
It references some kernels here:
https://github.com/pytorch/pytorch/blob/5d7770948587bb5f19b929200f20e2eaf7074e5c/aten/src/THCUNN/DistKLDivCriterion.cu 10
The CPU version is here:

  
      github.com
  
  
    pytorch/pytorch/blob/5d7770948587bb5f19b929200f20e2eaf7074e5c/aten/src/THNN/generic/DistKLDivCriterion.c 4
#ifndef TH_GENERIC_FILE
#define TH_GENERIC_FILE "generic/DistKLDivCriterion.c"
#else

void THNN_(DistKLDivCriterion_updateOutput)(
          THNNState *state,
          THTensor *input,
          THTensor *target,
          THTensor *output,
          bool sizeAverage,
          bool reduce)
{
  THNN_CHECK_SHAPE(input, target);

  if (!reduce) {
    THTensor_(resizeAs)(output, input);
    TH_TENSOR_APPLY3(real, input, real, target, real, output,
      *output_data = *target_data > 0 ? *target_data * (log(*target_data) - *input_data) : 0;
    );
    return;


  This file has been truncated. show original

  
  
    
    
  
  


Feel free to reply back if you have questions!