Prediction outputs differ for different batch sizes [BUG]

[melisande-c]

Describe the bug
If different batch sizes are used the predictions will be subtly different. This might be caused by PyTorch. In the example below I turned off tta_transforms to make sure that wasn't the cause.

To Reproduce
Code snippet allowing reproducing the behaviour:

import numpy as np
import matplotlib.pyplot as plt

from careamics import CAREamist
from careamics.config import create_n2v_configuration

config = create_n2v_configuration(
    experiment_name="PredBatchingTest", 
    data_type="array",
    axes="SYX",
    patch_size=[8, 8],
    batch_size=1,
    num_epochs=1,
    n_channels=1,
    )
image = np.random.random((1, 32, 32))

engine = CAREamist(source=config)
engine.train(train_source=image)
pred1 = engine.predict(
    source=image, 
    batch_size=1, 
    tile_size=(8, 8), 
    tile_overlap=(2,2), 
    tta_transforms=False
)
pred2 = engine.predict(
    source=image, 
    batch_size=2, 
    tile_size=(8, 8), 
    tile_overlap=(2,2), 
    tta_transforms=False
)
plt.imshow(abs(pred2-pred1))
plt.colorbar()

This produces the image:

As you can see the last tile produces the same output and this is because for this tile the batch size is equal to 1 for both predictions.

Additional context
I added this test at one point but it is skipped for now.

careamics/tests/test_careamist.py

Lines 588 to 623 in 0a29ea2

	@pytest.mark.skip(
	reason=(
	"This might be a problem at the PyTorch level during `forward`. Values up to "
	"0.001 different."
	)
	)
	def test_batched_prediction(tmp_path: Path, minimum_configuration: dict):
	"Compare outputs when a batch size of 1 or 2 is used"
	tile_size = (16, 16)
	tile_overlap = (4, 4)
	shape = (32, 32)
	train_array = random_array(shape)
	# create configuration
	config = Configuration(**minimum_configuration)
	config.training_config.num_epochs = 1
	config.data_config.axes = "YX"
	config.data_config.batch_size = 2
	config.data_config.data_type = SupportedData.ARRAY.value
	# instantiate CAREamist
	careamist = CAREamist(source=config, work_dir=tmp_path)
	# train CAREamist
	careamist.train(train_source=train_array)
	# predict with batch size 1 and batch size 2
	pred_bs_1 = careamist.predict(
	train_array, batch_size=1, tile_size=tile_size, tile_overlap=tile_overlap
	)
	pred_bs_2 = careamist.predict(
	train_array, batch_size=2, tile_size=tile_size, tile_overlap=tile_overlap
	)
	assert np.array_equal(pred_bs_1, pred_bs_2)