Check here for the entry point of Python custom functions once you call apply:
  

      github.com
  

  
    pytorch/pytorch/blob/f914ab193ebe5d432985af42a274dd21cb8c4fb8/torch/csrc/autograd/python_function.cpp#L581



      
      
            TORCH_CHECK(cdata,
              "Attribute 'name' is invalid for this instance of _C._FunctionBase. "
              "Accessing this attribute directly on an instance of autograd.Function is a legacy "
              "access pattern that is no longer supported. For examples on how to use new-style "
              "autograd functions, see "
              "https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function ");
            return THPUtils_packString(cdata->name());
            END_HANDLE_TH_ERRORS
          }
          
          
PyObject *THPFunction_apply(PyObject *cls, PyObject *inputs)
          {
            HANDLE_TH_ERRORS
            RECORD_FUNCTION(
              ((PyTypeObject*)cls)->tp_name,
              std::vector<c10::IValue>(),
              at::sequence_number::peek());
          
          
  THPObjectPtr backward_cls(PyObject_GetAttrString(cls, "_backward_cls"));
            if (!backward_cls) return nullptr;
            THPObjectPtr ctx_obj(PyObject_CallFunctionObjArgs(backward_cls, nullptr));
      
    
  


  

  
    
    
  

  


This function is where the grad_fn is set for the output. You can see that set_gradient_edge is called as long as the output is differentiable (i.e., its scalar type is floating or complex).
  

      github.com
  

  
    pytorch/pytorch/blob/5e72571df3f81a1ac9d0a2c96eff4acdac17579f/torch/csrc/autograd/custom_function.cpp#L28



      
      
          Variable VariableInfo::zeros(at::OptionalDeviceGuard& device_guard) const {
            if (is_empty) {
              // Return undefined tensor.
              return at::Tensor();
            } else {
              return at::zeros(
                  size, at::TensorOptions(scalar_type).device(device).layout(layout));
            }
          }
          
          
std::vector<c10::optional<Variable>> _wrap_outputs(const variable_list &input_vars,
            const std::unordered_set<at::TensorImpl*> &non_differentiable,
            const std::unordered_set<at::TensorImpl*> &dirty_inputs,
            const at::ArrayRef<c10::optional<Variable>> raw_outputs,
            const std::shared_ptr<Node> &cdata) {
          
          
  std::unordered_set<at::TensorImpl*> inputs;
            inputs.reserve(input_vars.size());
            for (auto& var : input_vars) {
              inputs.emplace(var.unsafeGetTensorImpl());
            }