The problem of the thickness variable being ineffective when using torch for optimization #566

I found that when using the** Torch backend** for optimization, when adding multiple thickness variables with problem.add_variable, usually only one is effective. And the other geometry properties remain unaffected. For example, in the "Optiland Torch Module - RMS Spot Size" case, if you also include the thickness of surface 1 and surface 2 in the optimization process, only the final added thickness will be valid.Is this due to my mistake or a bug?

import torch
import matplotlib.pyplot as plt
import optiland.backend as be
from optiland import optic, optimization
from optiland.ml import OpticalSystemModule

be.set_backend("torch")
be.grad_mode.enable()

lens = optic.Optic()
lens.surfaces.add(index=0, thickness=be.inf)
lens.surfaces.add(index=1, thickness=7, radius=1000, material="N-SF11", is_stop=True)
lens.surfaces.add(index=2, thickness=30, radius=-1000)
lens.surfaces.add(index=3)
lens.set_aperture(aperture_type="EPD", value=15)
lens.fields.set_type(field_type="angle")
lens.fields.add(y=0)
lens.wavelengths.add(value=0.55, is_primary=True)

problem = optimization.OptimizationProblem()
input_data = {
"optic": lens,
"surface_number": -1,
"Hx": 0, "Hy": 0,
"num_rays": 5,
"wavelength": 0.55,
"distribution": "hexapolar",
}
problem.add_operand("rms_spot_size", target=0, weight=1, input_data=input_data)
problem.add_variable(lens, "radius", surface_number=1)
problem.add_variable(lens, "thickness", surface_number=1)
problem.add_variable(lens, "radius", surface_number=2)
problem.add_variable(lens, "thickness", surface_number=2)

model = OpticalSystemModule(lens, problem)

params = list(model.parameters())
param_thick1 = params[1]
param_thick2 = params[3]

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

losses = []
grad_t1 = []
grad_t2 = []

for step in range(250):
optimizer.zero_grad()
loss = model()
loss.backward()

g1 = param_thick1.grad
g2 = param_thick2.grad
grad_t1.append(g1.item() if g1 is not None else 0.0)
grad_t2.append(g2.item() if g2 is not None else 0.0)

optimizer.step()
model.apply_bounds()
losses.append(loss.item())
problem.info()

plt.figure()
plt.plot(losses)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("Lens Optimization: Spot Size Reduction")
plt.grid(alpha=0.25)
plt.show()

plt.figure()
plt.plot(grad_t1, label="Thickness 1 gradient")
plt.plot(grad_t2, label="Thickness 2 gradient")
plt.xlabel("Iteration")
plt.ylabel("Gradient Value")
plt.title("Gradient History of Thickness Variables")
plt.legend()
plt.grid(alpha=0.25)
plt.show()

_ = lens.draw()