AttributeError: 'SegmentConsensusBackward' object has no attribute 'consensus_type'

I have this error appears after i have added . apply to object class calling as it was written this function has been deprecated when i used torch.autograde.function
the following is the code:
import torch

import math

import numpy

class Identity(torch.nn.Module):

def forward(self, input):

    return input

class SegmentConsensus(torch.autograd.Function):

def __init__(self, consensus_type, dim=1):

    self.consensus_type = consensus_type

    self.dim = dim

    self.shape = None

def forward(self, input_tensor):



    #self.shape = len(input_tensor)


    if self.consensus_type.apply == 'avg':

        output = input_tensor.mean(dim=self.dim, keepdim=True)

    elif self.consensus_type == 'identity':

        output = input_tensor


        output = None

    return output

def backward(self, grad_output):

    if self.consensus_type == 'avg':

        grad_in = grad_output.expand(self.shape) / float(self.shape[self.dim])

    elif self.consensus_type == 'identity':

        grad_in = grad_output


        grad_in = None

    return grad_in

class ConsensusModule(torch.nn.Module):

def __init__(self, consensus_type, dim=1):

    super(ConsensusModule, self).__init__()

    self.consensus_type = consensus_type if consensus_type != 'rnn' else 'identity'

    self.dim = dim

def forward(self, input):

    return SegmentConsensus((self.consensus_type, self.dim)).apply(input)

and this is the error
f self.consensus_type.apply == ‘avg’:
AttributeError: ‘SegmentConsensusBackward’ object has no attribute ‘consensus_type’

I guess the error might be raised, since you are using the legacy definition of custom autograd.Functions and would have to adapt the code to the new approach as explained here.