Add missing return in example code
Add missing return in example code

Keras custom layers example code missed a return statement, which
caused layers to not have a config when the example was followed.
tensorflow/docs,tensorflow/docs,tensorflow/docs,tensorflow/docs-l10n