It looks like input to NN should have [batch_size, width, height] dimensions instead of [batch_size, 1, width, height], but I'm not completely sure. In case of conv and mlp input is reshaped explicitly with Reshape layer, in case of linear it's most likely reshaped implicitly to [batch_size, width*height]. Still it's not very clear what that 1 is for, looks like leftover from conv...
It looks like input to NN should have
[batch_size, width, height]dimensions instead of[batch_size, 1, width, height], but I'm not completely sure. In case ofconvandmlpinput is reshaped explicitly withReshapelayer, in case oflinearit's most likely reshaped implicitly to[batch_size, width*height]. Still it's not very clear what that1is for, looks like leftover fromconv...