supposedly fix cuda for IsoRun

axsk · Feb 14, 2024 · e16e90e · e16e90e
1 parent c9d9543
commit e16e90e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/cuda.jl b/src/cuda.jl
@@ -9,6 +9,7 @@ move the model and data of the given `IsoRun` to the GPU for CUDA support
 function gpu!(iso::IsoRun)
     iso.model = Flux.gpu(iso.model)
     iso.data = Flux.gpu(iso.data)
+    iso.opt = Flux.gpu(iso.opt)
     return
 end