You could check for axis=None before bailing out since in that case you would insert a GpuFlatten anyway.