PyTorch Backend Error

[muslikhah10-alt]

Code：
import optiland.backend as be
be.set_backend("torch") # Set the backend to PyTorch
be.set_precision("float32") # or "float64" for double precision
be.grad_mode.enable() # Enable gradient tracking
be.set_device("cuda") # run on GPU if available, otherwise use "cpu" option

from optiland.samples.objectives import CookeTriplet
from optiland.psf import FFTPSF

lens = CookeTriplet()
psf = FFTPSF(lens, field=(0, 1), wavelength=0.55)
psf.view(projection="2d", num_points=256)

Error：
Traceback (most recent call last):
File "C:\Users\Administrator\PycharmProjects\PythonProject\start.py", line 11, in
psf = FFTPSF(lens, field=(0, 1), wavelength=0.55)
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\psf\fft.py", line 110, in init
super().init(
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\psf\base.py", line 87, in init
super().init(
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\wavefront\wavefront.py", line 82, in init
self._generate_data()
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\wavefront\wavefront.py", line 162, in _generate_data
data = self.strategy.compute_wavefront_data(field, wl)
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\wavefront\strategy.py", line 207, in compute_wavefront_data
rays = self.optic.trace(*field, wavelength, None, self.distribution)
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\optic\optic.py", line 609, in trace
return self.ray_tracer.trace(Hx, Hy, wavelength, num_rays, distribution)
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\raytrace\real_ray_tracer.py", line 79, in trace
rays = self.ray_generator.generate_rays(
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\rays\ray_generator.py", line 52, in generate_rays
x0, y0, z0 = self.optic.field_definition.get_ray_origins(
File "C:\ProgramData\anaconda3\envs\pytorch-gpu\lib\site-packages\optiland\fields\field_types.py", line 164, in get_ray_origins
x = -be.tan(be.radians(field_x)) * (offset + EPL)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

What's the reason? How to fix it?

[anntnzrb]

Initialize the lens AFTER setting backend/device. The issue at field_types.py:164 occurs because be.set_device("cuda") doesn't transfer existing data. Call be.set_backend("torch") and be.set_device("cuda") before creating any optical objects (like CookeTriplet()) to ensure all tensors are allocated on the correct device from the start.

[wuxujiong]

I did this Call be.set_backend("torch") and be.set_device("cuda") before creating any optical objects (like CookeTriplet()), but still raise the same problem

[muslikhah10-alt]

I’ve already fixed this issue! You can take a look at #426