Extract feature from a model generated using Engine

I am trying to extract feature using the code of https://github.com/szagoruyko/wide-residual-networks.
I have an Engine model. I have model which looks like:
def f(input, params, mode):
x = F.conv2d(input, params[‘conv0’], padding=1)
g0 = group(x, params, ‘group0’, mode, 1)
g1 = group(g0, params, ‘group1’, mode, 2)
g2 = group(g1, params, ‘group2’, mode, 2)
o = F.relu(utils.batch_norm(g2, params, ‘bn’, mode))
o = F.avg_pool2d(o, 8, 1, 0)
o = o.view(o.size(0), -1)
o = F.linear(o, params[‘fc.weight’], params[‘fc.bias’])
return o
It is a function instead of a class.

Lets assume I save a trained model as torch.save(dict( params=params, epoch=t[‘epoch’], optimizer=state[‘optimizer’].state_dict()), path ). I want to reuse this model for extracting features later.

Since it is not a class, I can not create an object out of it and initialize with pretrained parameters. So my solution is to change the function definition (f_1) above(discard last two layers). Then call
f_1, params = resnet(opt.depth, opt.width, num_classes)
state_dict = torch.load(path_to_saved_model)
params_tensors = state_dict[‘params’]
for k, v in params.items():

This initializes the model f_1 with the subset layers parameter from pretrained model.
I am okay upto this. But then I have no idea how do I get prediction(features) out of this function (with Engine included in this model , it becomes more difficult for me). Just to clarify again: I want features from a pretrained model. Sorry about long post and maybe this may be a naive post to experts here.
The basic training model for Engine is here : https://github.com/szagoruyko/wide-residual-networks/blob/master/pytorch/main.py