Thank you for your quick response. I am trying to save on the computation time of a layer, which has been detached from the back propagation.
The original code is from the paper
FreezeOut: Accelerate Training by Progressively Freezing Layers (https://arxiv.org/abs/1706.04983 7)
My modified code is given below
class Bottleneck(nn.Module):
    def __init__(self, nChannels, growthRate,layer_index, train_size, test_size, batch_sz):
        super(Bottleneck, self).__init__()
        interChannels = 4*growthRate
        self.bn1 = nn.BatchNorm2d(nChannels)
        self.conv1 = nn.Conv2d(nChannels, interChannels, kernel_size=1,
                               bias=False)
        self.bn2 = nn.BatchNorm2d(interChannels)
        self.conv2 = nn.Conv2d(interChannels, growthRate, kernel_size=3,
                               padding=1, bias=False)

        # If the layer is still being trained
        self.active=True
        
        # The index of this layer relative to the overall net
        self.layer_index=layer_index
        self.save_out = False

#Counter to keep track of which batch is being processed
        self.counter = 0
        self.train_size = train_size
        self.test_size = test_size
        self.batch_sz = batch_sz
 
        remainder = train_size % batch_sz

        self.maxCounter = train_size//batch_sz


        if(remainder != 0):
            self.maxCounter += 1
			
# List to hold the output/detach of inactive layers            
        self.out = [None] * (self.maxCounter)

        
    def forward(self, x):
    
        
        # If we're not active, return a detached output to prevent backprop.
        if self.active: 
            out = self.conv1(F.relu(self.bn1(x)))
            out = self.conv2(F.relu(self.bn2(out)))
            out = torch.cat((x, out), 1)
            self.counter+=1
            if(self.counter >= self.maxCounter):
                self.counter = 0      
            return out

        else:
# Layer has become inactive, check if we have already the output/detach for the current input, if not store it         
            if (self.out[self.counter] is not None):
                detach = self.out[self.counter]
                self.counter+=1
                if(self.counter >= self.maxCounter):
                    self.counter = 0
					
				detach.volatile = False	
                return detach
                 
            else:
                out = self.conv1(F.relu(self.bn1(x)))
                out = self.conv2(F.relu(self.bn2(out)))
                out = torch.cat((x, out), 1)
# In place detach 			
                out.detach_()

                self.out[self.counter] = out
                self.counter+=1
                if(self.counter >= self.maxCounter):
                    self.counter = 0            

                return out