Collecting features from network.foward()

So basically I want to achieve the same goal as in this code but in TensorFlow

def get_function(network, loader):
    ''' Collect function (features) from the routine '''
    features = []
    for batch_idx, (inputs, targets) in enumerate(loader):
        inputs, targets ='cpu'),'cpu')

        features.append([f.cpu().data.numpy().astype(np.float16) for f in network.forward_features(inputs)])

    return [np.concatenate(list(zip(*features))[i]) for i in range(len(features[0]))]

Are there any clean ways to do this with TensorFlow iterator?

Reference for future readers.