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!