Change in Loss function (classification -> regression) causes MXNet reshape error

[hmf]

I am trying to get a simple Autoencoder to work with the MNIST data (see #601).
If I use that very simple network with this classification loss:

    DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
      .optDevices(trainSetUp.devices)
      .addTrainingListeners(TrainingListener.Defaults.logging(outputDir):_*)
      .addTrainingListeners(listener)
      .optOptimizer( Adam.builder().build )

training proceeds with EasyTrain without a problem (even though the result is no what I need). In order to make this a regression problem, I tried the L2 and the Sigmoid Binary Cross Entropy losses so:

    DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())

and

    DefaultTrainingConfig(Loss.l2Loss())

In both cases I get the error:

Failure(ai.djl.engine.EngineException: MXNet engine call failed: MXNetError: Check failed: src.Size() == dst->Size() (256 vs. 200704) : Cannot reshape array of size 256 into shape [256,784]
Stack trace:
  File "src/operator/numpy/np_matrix_op.cc", line 145
)
MXNet engine call failed: MXNetError: Check failed: src.Size() == dst->Size() (256 vs. 200704) : Cannot reshape array of size 256 into shape [256,784]
Stack trace:
  File "src/operator/numpy/np_matrix_op.cc", line 145

I note that 200704 = 256 * 784 = 256 * 28 * 28. This seems strange to me because both input and labels have the following shape:

images.getShape = (256, 1, 28, 28)
labels.getShape = (256, 1, 28, 28)

I was expecting the shape [256,784] and not a single 200704 array (matrix?). I also assumed that the regression losses would be a simple drop-in replacement because I have not changed any of the input pre-processing. I do have a val trainer: Trainer that is set with:

        trainer.initialize(trainSetUp.inputShape)

How are those loss functions different? What must be changed or set in these cases? How should one go about diagnosing this? For reference I show below the NN used in this case.

TIA

  def net1(): SequentialBlock = 
    val net = SequentialBlock()
    net.add(Blocks.batchFlattenBlock(784))
    net.add(Linear.builder().setUnits(32).build())
    //net.add(Activation::relu)
    net.add(relu)
    net.add(Linear.builder().setUnits(784).build())
    net.add(sigmoid)
    net.setInitializer( NormalInitializer(0.01f) ) 
    net

[hmf]

Finally figured it out. I was using the incorrect data set. I was using the dataset that uses a vector of integers as a label (classes).

Apologies for the noise.