Improper convergence for a multidimensional simulator for SNRE and SNLE

[paarth-dudani]

Describe the bug
I don't get a proper convergence with SNLE and SNRE for a multidimensional prior and simulatoe with noise.

To Reproduce

1. Versions
   Python version: 3.9.13
   SBI version: 0.23.1

2. Minimal code example

scale_true = torch.tensor([[2.0, 1.5, 1.0, 2.0]])

# Define the prior
num_dim = 4
num_sim = 1000
num_rounds = 2
initial_cond = torch.tensor([1.0, 1.0, 2.0,1.2])


n = 1000
t = np.linspace(0,10,n)

# prior

uni_prior_2D = BoxUniform(low=0* torch.ones(num_dim), high=3* torch.ones(num_dim))

# Simulators

def simulator_exp(input, theta, num_sim, n_dim):
    x = np.zeros((n_dim, num_sim, len(input)))
    for i in range(n_dim):
            x[i] = (1/initial_cond[i])*torch.exp(-input/theta[:,i].reshape(-1,1)) + 0.5*torch.rand_like(theta[:,i]).reshape(-1,1)
    return torch.tensor(x).reshape(num_sim, n_dim, len(input)).float()

x_o = simulator_exp(t, scale_true, 1, num_dim)
x_o = x_o.reshape(1, num_dim*len(t))


inference = NLE(uni_prior_2D)
proposal = uni_prior_2D
for _ in range(num_rounds):
    theta = proposal.sample((num_sim,))
    x = simulator_exp(t, theta, num_sim, num_dim)  # Simulate data
    x = x.reshape(num_sim, num_dim*len(t))
    x = np.clip(x, 0, 8)
    _ = inference.append_simulations(theta, x).train()
    posterior = inference.build_posterior(mcmc_method="slice_np_vectorized",
                                          mcmc_parameters={"num_chains": 5,
                                                           "thin": 5})
    proposal = posterior.set_default_x(x_o)

All else same for SNLE.

3. Error: None. Just poor convergence.

SNRE:

SNLE:

Expected behavior
A sharper and proper convergence of the posterior