can't use lag method

[2M-kotb]

i need to formalize my forecasting problem as predicting one-step-ahead using lag say equal 3.

when i do that the patch in will be [num of patches, patch_size, 3, 1] and patch out will be [num of patches, patch_size, 1, 1] and that gives me the following error in train method

InvalidArgumentError (see above for traceback): ConcatOp : Dimensions of inputs should match: shape[0] = [3,1,25] vs. shape[1] = [1,1,1]