Currently, my last bit of my model looks like this:
  (conv_last): Sequential(
    (0): Conv2d(2688, 5376, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): BatchNorm2d(5376, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace)
  )
  (globalpool): Sequential(
    (0): AvgPool2d(kernel_size=7, stride=7, padding=0)
  )
  (classifier): Sequential(
    (0): Linear(in_features=5376, out_features=2300, bias=True)
  )

with forward function:
    def forward(self, x):
        x = self.conv1(x)
        x = self.features(x)
        x = self.conv_last(x)
        x = self.globalpool(x)
        x = x.view(-1, self.stage_out_channels[-1])
        closs_output = self.linear_closs(x)
        x = self.classifier(x)
        return x

I wish to add a centre loss function, where the end result is like:

  (conv_last): Sequential(
    (0): Conv2d(2688, 5376, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): BatchNorm2d(5376, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace)
  )
  (globalpool): Sequential(
    (0): AvgPool2d(kernel_size=7, stride=7, padding=0)
  )
  (classifier): Sequential(
    (0): Linear(in_features=5376, out_features=2300, bias=True)
  )
  (linear_closs): Linear(in_features=5376, out_features=1280, bias=False)
)

with new forward function:
    def forward(self, x):
        x = self.conv1(x)
        x = self.features(x)
        x = self.conv_last(x)
        x = self.globalpool(x)
        x = x.view(-1, self.stage_out_channels[-1])
        closs_output = self.linear_closs(x)
        x = self.classifier(x)
        return closs_output, x
#         return x

I used torch.load(model). Is there a way to add these new lines into the model? The only changes are the last layer of the model, and 2 new lines in the forward function to specify the return for the new loss function.
Thanks!