store only features, not the coordinates on gpu

axsk · May 11, 2024 · 5f36c01 · 5f36c01
1 parent 86740e1
commit 5f36c01
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/simulation.jl b/src/simulation.jl
@@ -90,8 +90,8 @@ end
 
 #features(sim::SimulationData, x) = sim.featurizer(x)
 
-gpu(d::SimulationData) = SimulationData(d.sim, gpu(d.features), gpu(d.coords), d.featurizer)
-cpu(d::SimulationData) = SimulationData(d.sim, cpu(d.features), cpu(d.coords), d.featurizer)
+gpu(d::SimulationData) = SimulationData(d.sim, gpu(d.features), d.coords, d.featurizer)
+cpu(d::SimulationData) = SimulationData(d.sim, cpu(d.features), d.coords, d.featurizer)
 
 features(d::SimulationData, x) = d.featurizer(x)