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 self.network.module.forward_features() routine '''
features = []
for batch_idx, (inputs, targets) in enumerate(loader):
inputs, targets = inputs.to('cpu'), targets.to('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?