non GPU devices

vpuri3 · May 13, 2024 · 777cbbe · 777cbbe
1 parent 907a038
commit 777cbbe
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/examples/eg1.jl b/examples/eg1.jl
@@ -43,8 +43,13 @@ function main()
 
     stM, stK = device(stM), device(stK)
 
-    @btime CUDA.@sync $mlp($x, $pM, $stM)
-    @btime CUDA.@sync $kan($x, $pK, $stK)
+    if device isa LuxDeviceUtils.AbstractLuxGPUDevice
+        @btime CUDA.@sync $mlp($x, $pM, $stM)
+        @btime CUDA.@sync $kan($x, $pK, $stK)
+    else
+        @btime $mlp($x, $pM, $stM)
+        @btime $kan($x, $pK, $stK)
+    end
 
     nothing
 end