Could you please let me know where am I making mistake?
ENet = EfficientNet.from_pretrained('efficientnet-b5')

class ENET(nn.Module):
  def __init__(self):
    super(ENET, self).__init__()
    
    self.ENet = ENet

    #disect the network to access its last convolutional layer
    self.head0 = nn.Sequential(*list(self.ENet.children())[:4])
    self.BN =  nn.BatchNorm2d(2048, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
    self.tail = nn.Sequential(*list(self.ENet.children())[-4:])

    # placeholder for the gradients
    self.gradients = None

# hook for the gradients of the activations
def activations_hook(self, grad):
    self.gradients = grad
    
def forward(self, x):
    self.x, self.labels = self.head0(x)
    #register the hook
    h = self.x.register_hook(self.activations_hook)       
    x = self.BN(self.x)
    x = x.view((1, -1))
    x = self.tail(x)
    return x

# method for the gradient extraction
def get_activations_gradient(self):
    return self.gradients

# method for the activation exctraction
def get_activations(self, x):
    return self.head0(x)


model = ENET()
# set the evaluation mode
model.eval()
# get the image from the dataloader
 img, _= next(iter(train_loader))
 img = img.to(device)
 print(img.shape)
# get the most likely prediction of the model
pred = model(self.X)

error:
torch.Size([1, 3, 224, 224])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-46c7c5e707f6> in <module>()
  9 print(img.shape)
 10 # get the most likely prediction of the model
 ---> 11 pred = model(img)

 4 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, 
**kwargs)
 725             result = self._slow_forward(*input, **kwargs)
726         else:
--> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

 <ipython-input-7-4bf4b58e1763> in forward(self, x)
 26 
 27     def forward(self, x):
 ---> 28         self.x, self.labels = self.head0(x)
 29 
 30         #x = self.head(x)

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, 
**kwargs)
725             result = self._slow_forward(*input, **kwargs)
726         else:
 --> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/container.py in forward(self, input)
115     def forward(self, input):
116         for module in self:
--> 117             input = module(input)
118         return input
119 

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, 
**kwargs)
725             result = self._slow_forward(*input, **kwargs)
726         else:
--> 727             result = self.forward(*input, **kwargs)
728         for hook in itertools.chain(
729                 _global_forward_hooks.values(),

TypeError: forward() takes 1 positional argument but 2 were given