I have a reshaping problem that I don\'t seem to get solved. For this I have a following code available for my prototypical network:
class PrototypicalNetwork