Cannot implement customized loss function?

Could you please help me to understand, how do i implement a custom loss?
I do as following and its wan’t work.

import mxnet as mx
import logging
logging.basicConfig(level='DEBUG')

use_custom = False

A =  mx.nd.random.uniform(-1, 1, (5, 1))
B =  mx.nd.random.uniform(-1, 1)

X = mx.nd.random.uniform(-1, 1, (100, 5))
y = mx.nd.dot(X, A) + B

iter = mx.io.NDArrayIter(data=X, label=y, data_name='data', label_name='label', batch_size=20, shuffle=True)

data  = mx.sym.Variable('data')
label = mx.sym.Variable('label')

net = mx.sym.FullyConnected(data, num_hidden=1)
if use_custom:
    net = mx.sym.MakeLoss(mx.sym.square(net - label))
else:
     net = mx.sym.LinearRegressionOutput(net, label=label)

mod = mx.mod.Module(net, label_names=('label',))
mod.fit(iter, num_epoch=50, eval_metric='mse', optimizer='adam')

thank you.

Thanks for the easy to reproduce example :slight_smile:
edit: I had forgotten to flip the use_custom flag. As pointed out by @safrooze your custom loss is working however the evaluation metric is using the output of your loss and comparing it with the label. Hence it looks like it is not converging but in fact it is.

mod = mx.mod.Module(net, label_names=['label'])
identity = mx.metric.CustomMetric(lambda x,y:y, name='mse_id')
mod.fit(iter, num_epoch=10, eval_metric=identity, optimizer='adam')

This gives you this:

INFO:root:Epoch[0] Train-mse_id=0.434285
INFO:root:Epoch[0] Time cost=0.056
INFO:root:Epoch[1] Train-mse_id=0.000387
INFO:root:Epoch[1] Time cost=0.055
INFO:root:Epoch[2] Train-mse_id=0.000000
INFO:root:Epoch[2] Time cost=0.055
INFO:root:Epoch[3] Train-mse_id=0.000000
INFO:root:Epoch[3] Time cost=0.055
INFO:root:Epoch[4] Train-mse_id=0.000000
INFO:root:Epoch[4] Time cost=0.055
INFO:root:Epoch[5] Train-mse_id=0.000000
INFO:root:Epoch[5] Time cost=0.056
INFO:root:Epoch[6] Train-mse_id=0.000000
INFO:root:Epoch[6] Time cost=0.056
INFO:root:Epoch[7] Train-mse_id=0.000000
INFO:root:Epoch[7] Time cost=0.056
INFO:root:Epoch[8] Train-mse_id=0.000000
INFO:root:Epoch[8] Time cost=0.056
INFO:root:Epoch[9] Train-mse_id=0.000000
INFO:root:Epoch[9] Time cost=0.056

Thank you,
After giving a relevant metric, it works.