Sure. I will use the custom leaky relu as example for ease. Also i am using conda. pytorch version 1.7.1
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Function # import Function to create custom activations


import numpy.matlib

from torch.autograd import gradcheck

import random
from torch.autograd import Variable

torch.manual_seed(0)
torch.cuda.manual_seed(0)
np.random.seed(0)
random.seed(0)
torch.backends.cudnn.deterministic=True

CUDA_LAUNCH_BLOCKING=1
# Device configuration
#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device("cuda:0") 
#device = 'cpu'
#torch.backends.cuda.matmul.allow_tf32 = False 
# Hyper-parameters 
num_epochs = 5
batch_size = 32
learning_rate = 0.001
L_relu_lambda = 0.02



transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5), (0.5))])

train_dataset = torchvision.datasets.FashionMNIST(root='./data', train= True,
                                        download=True, transform=transform)

test_dataset = torchvision.datasets.FashionMNIST(root='./data', train= False,
                                       download=True, transform=transform)


train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64,
                                          shuffle=True, worker_init_fn=np.random.seed(12))

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64,
                                         shuffle=False, worker_init_fn=np.random.seed(12))






def train_model(model,trainloader):
    '''
    Function trains the model and prints out the training log.
    '''
    #setup training
    
    #define loss function
    criterion = nn.NLLLoss()
    #define learning rate
    learning_rate = 0.003
    #define number of epochs
    epochs = 5
    #initialize optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    #run training and print out the loss to make sure that we are actually fitting to the training set
    print('Training the model. Make sure that loss decreases after each epoch.\n')
    for e in range(epochs):
        running_loss = 0
        for images, labels in trainloader:
            images = images.to(device)
            labels = labels.to(device)
            images = images.view(images.shape[0], -1)
            log_ps = model(images)
            loss = criterion(log_ps, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
        else:
            # print out the loss to make sure it is decreasing
            print(f"Training loss: {running_loss}")


class Lrelu_custom(Function):
    
    #both forward and backward are @staticmethods
    @staticmethod
    def forward(ctx, input,alpha):
        
        ctx.save_for_backward(input) # save input for backward pass
        output= input.clone()
        output= Variable(input.clamp(min=0)+input.clamp(max=0)*alpha, requires_grad = True)
       
        ctx.slope = alpha
        
        return output
        

    @staticmethod
    def backward(ctx, grad_output):
        
        grad_input = None # set output to None

        input, = ctx.saved_tensors # restore input from context
        slope = ctx.slope
        grad_input = grad_output.clone()
        grad_input = grad_input * (input > 0).float() + grad_input * (input < 0).float() * slope
        
        
        return grad_input, None
    
class Classifierleakyrelu(nn.Module):
    '''
    Simple fully-connected classifier model to demonstrate leakyrelu activation.
    '''
    def __init__(self):
        super(Classifierleakyrelu, self).__init__()

        # initialize layers
        self.fc1 = nn.Linear(784, 256)
        
        # create shortcuts for leaky relu
        self.a1 = Lrelu_custom.apply
        
        
    def forward(self, x):
        # make sure the input tensor is flattened
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        # apply leaky relu
        x = self.a1(x,0.001)
       
        x = F.log_softmax(x, dim=1)
        
        return x
    
model = Classifierleakyrelu().to(device)
train_model(model,train_loader)

cpu results and gpu results varied on my custom code.so i just tested on this simple leaky relu and check if the code ever accessed the custom backward during back pass, used breakpoints in spyder for it but it never stopped at the break points. Thanks again