Solved by ptrblck in post #2 
              
              
                The error is raised, because you are trying to feed the output tensor again to the complete model via self.features. 
Remove self.features as it’s not doing anything at the moment. 
Once this bug is fixed, you’ll run into a shape mismatch, as our activation won’t have 128 features to be suitable for…