Custom Keras Layer Troubles

Hate to ask a question like this on machine learning but googling has yielded nothing useful - I've just found 2 github threads where people on super old versions of tensorflow got the same error message, but not for the same reason im getting it.

Basically; I'm implementing this facial point paper for work; and it uses spatial softargmax (just a layer that takes in a stack of images a lot like this - and it returns the most "intense part" of the image (so just the x,y coordinates of the white blob). It takes in an array with 68 of these images (all 1 channel, so the array is 100x100x68) and it gives 68 pairs of x,y coordinates for each one - these end up being the facial points.

The layer I have written in keras to do this is;

class spatial_softArgmax(Layer):

        def __init__(self, output_dim, **kwargs):
                self.output_dim = output_dim
                super(spatial_softArgmax, self).__init__(**kwargs)


        def build(self, input_shape):
                # Create a trainable weight variable for this layer.
                # self.kernel = self.add_weight(name='kernel', shape=(input_shape[1], self.output_dim), initializer='uniform', trainable=True)
                super(spatial_softArgmax, self).build(input_shape)  # Be sure to call this somewhere!

        def call(self, x):
            #print "input to ssm:  " + str(x.shape)
            maps = layers.spatial_softmax(x, data_format='NCHW')
            #maps = tf.reshape(maps, [2, 68])
            #print "output of ssm: " + str(maps.shape)
            return maps

        def get_config(self):
                config = super(spatial_softArgmax, self).get_config()
                config['output_dim'] = self.output_dim
                config['input_shape'] = self.input_shape
                return config

        def compute_output_shape(self, input_shape):
                return (input_shape[0], self.output_dim)

It isn't working though; like at all; I keep getting this error that says 'None Values not supported' whenever i try to start training (training begins if i replace it with a dense - so the problem is definitely this layer) - which makes me think that my layer isn't returning anything?? I dug around in the TF code a bit around where this exception is raise but didn't really find much...

If you guys see anything wrong with my layer just by glancing i'd really appreciate it if you'd let me know; I need to have this network training overnight tonight.

EDIT: I've removed the self.kernel line; but am now getting:

x,y shape (12000, 3, 100, 100) - (12000, 68, 2)
Traceback (most recent call last):
  File "architecture.py", line 58, in <module>
    model.fit(x, y, batch_size=1, epochs=50, verbose=1)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/training.py", line 1522, in fit
    batch_size=batch_size)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/training.py", line 1382, in _standardize_user_data
    exception_prefix='target')
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/training.py", line 132, in _standardize_input_data
    str(array.shape))
ValueError: Error when checking target: expected spatial_soft_argmax_1 to have 2 dimensions, but got array with shape (12000, 68, 2)

You've added a weight self.kernel but didn't use it anywhere in your call function.

Therefore, TF is unable to compute gradient for the parameters in self.kernel. In this case, the gradient will be None and that's why you're seeing an error complaining about operating on None values.

Removing the line self.kernel = self.add_weight(...) should make it work.