Injecting input data in the output layer

Question

I'm building a model using Tensorflow where the input is a slice of the output. Think of the output layer as a 2D array. One row of that array is the input data. The neural network currently tries to connect the input to the output using a mean-squared error loss function. It's doing a fairly good job, but the accuracy needs to be improved a little.

To do that, I'm trying to add another physics-based loss function. If I can have the network place the input slice in its correct location in the output, that would greatly simplify the problem as each row in the output 2D array depends on the two rows above it.

I hope this makes sense.