diff --git a/bench/comparison.md b/bench/comparison.md
index 2343c64..f6c9830 100644
--- a/bench/comparison.md
+++ b/bench/comparison.md
@@ -1,41 +1,49 @@
-## NeuralOperators.jl (Lux)
+# NeuralOperators.jl Benchmarks
 
-## FNO
+## Fourier Neural Operators
 
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-| 1 | 14.173699999999998 ms | 755.1466 ms |
-| 2 | 29.118399999999998 ms | 1407.2298 ms |
-| 3 | 37.6924 ms | 2367.5004999999996 ms |
-| 4 | 41.431400000000004 ms | 3035.1971 ms |
-| 5 | 59.305 ms | 3456.1902999999998 ms |
+### Lux.jl (Julia)
 
-## FNO (python: neuraloperator)
+| #layers | Forward               | Train: 10 epochs      |
+|:------- |:--------------------- |:--------------------- |
+| 1       | 14.173699999999998 ms | 755.1466 ms           |
+| 2       | 29.118399999999998 ms | 1407.2298 ms          |
+| 3       | 37.6924 ms            | 2367.5004999999996 ms |
+| 4       | 41.431400000000004 ms | 3035.1971 ms          |
+| 5       | 59.305 ms             | 3456.1902999999998 ms |
 
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-|  1  |  5.731542900000932  ms |  17.667421199992532  ms |
-|  2  |  7.833489999989979  ms |  25.585920999990776  ms |
-|  3  |  10.18306370000937  ms |  33.69801080002799  ms |
-|  4  |  12.33892210002523  ms |  41.98180860001594  ms |
-|  5  |  14.732645300013246  ms |  50.13744520000182  ms |
+### Flux.jl (Julia)
+
+### neuraloperator (Python)
+
+| #layers | Forward                | Train: 10 epochs       |
+|:------- |:---------------------- |:---------------------- |
+| 1       | 5.731542900000932  ms  | 17.667421199992532  ms |
+| 2       | 7.833489999989979  ms  | 25.585920999990776  ms |
+| 3       | 10.18306370000937  ms  | 33.69801080002799  ms  |
+| 4       | 12.33892210002523  ms  | 41.98180860001594  ms  |
+| 5       | 14.732645300013246  ms | 50.13744520000182  ms  |
 
 ## DeepONet
 
-| #layers | Forward | Train: 10 epochs | 
-| --- | --- | --- | 
-| 1 | 3.3952750000000003 ms | 76.604576 ms | 
-| 2 | 4.360458 ms | 104.460251 ms | 
-| 3 | 5.6310780000000005 ms | 149.148633 ms | 
-| 4 | 7.199777 ms | 178.464657 ms | 
-| 5 | 7.8226819999999995 ms | 193.760173 ms |
-
-## DeepONet (python: deepxde)
-
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-|  1  |  0.7689221948385239  ms |  25.76469287276268  ms |
-|  2  |  0.7733150571584702  ms |  32.17746138572693  ms |
-|  3  |  0.8474267274141312  ms |  36.93301998078823  ms |
-|  4  |  1.0069304704666138  ms |  45.45578710734844  ms |
-|  5  |  1.406572386622429  ms |  59.06449243426323  ms |
+### Lux.jl (Julia)
+
+| #layers | Forward               | Train: 10 epochs |
+|:------- |:--------------------- |:---------------- |
+| 1       | 3.3952750000000003 ms | 76.604576 ms     |
+| 2       | 4.360458 ms           | 104.460251 ms    |
+| 3       | 5.6310780000000005 ms | 149.148633 ms    |
+| 4       | 7.199777 ms           | 178.464657 ms    |
+| 5       | 7.8226819999999995 ms | 193.760173 ms    |
+
+### Flux.jl (Julia)
+
+### deepxde (Python)
+
+| #layers | Forward                | Train: 10 epochs      |
+|:------- |:---------------------- |:--------------------- |
+| 1       | 0.7689221948385239  ms | 25.76469287276268  ms |
+| 2       | 0.7733150571584702  ms | 32.17746138572693  ms |
+| 3       | 0.8474267274141312  ms | 36.93301998078823  ms |
+| 4       | 1.0069304704666138  ms | 45.45578710734844  ms |
+| 5       | 1.406572386622429  ms  | 59.06449243426323  ms |
diff --git a/bench/flux.jl b/bench/flux.jl
new file mode 100644
index 0000000..30580d7
--- /dev/null
+++ b/bench/flux.jl
@@ -0,0 +1,66 @@
+using ThreadPinning
+pinthreads(:cores)
+threadinfo()
+
+using BenchmarkTools, NeuralOperators, Random, Optimisers, Zygote
+
+# TODO: Add training code
+
+# FNO
+n_points = 128
+batch_size = 64
+
+x = rand(Float32, 1, n_points, batch_size);
+y = rand(Float32, 1, n_points, batch_size);
+data = [(x, y)];
+t_fwd = zeros(5)
+t_train = zeros(5)
+
+for i in 1:5
+    chs = (1, 128, fill(64, i)..., 128, 1)
+    model = FourierNeuralOperator(; ch=chs, modes=(16,), σ=gelu)
+    # model(x) # TTFX
+
+    # t_fwd[i] = @belapsed $model($x)
+
+    # t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
+end
+
+println("\n## FNO (Flux NeuralOperators.jl)")
+print("| #layers | Forward | Train: 10 epochs | \n")
+print("| --- | --- | --- | \n")
+for i in 1:5
+    print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
+end
+
+# DeepONets
+eval_points = 128
+batch_size = 64
+dim_y = 1
+m = 32
+
+u = rand(Float32, m, batch_size);
+y = rand(Float32, dim_y, eval_points, batch_size);
+
+g = rand(Float32, eval_points, batch_size);
+
+data = [((u, y), g)]
+t_fwd = zeros(5)
+t_train = zeros(5)
+for i in 1:5
+    ch_branch = (m, fill(64, i)..., 128)
+    ch_trunk = (dim_y, fill(64, i)..., 128)
+    model = DeepONet(ch_branch, ch_trunk)
+    # model(u, y) # TTFX
+
+    # t_fwd[i] = @belapsed $model($u, $y)
+
+    # t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
+end
+
+println("\n## DeepONet (Flux NeuralOperators.jl)")
+print("| #layers | Forward | Train: 10 epochs | \n")
+print("| --- | --- | --- | \n")
+for i in 1:5
+    print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
+end
diff --git a/bench/flux/Project.toml b/bench/flux/Project.toml
new file mode 100644
index 0000000..d6d94c9
--- /dev/null
+++ b/bench/flux/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+ThreadPinning = "811555cd-349b-4f26-b7bc-1f208b848042"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/bench/lux.jl b/bench/lux.jl
index ad2492b..e48d73c 100644
--- a/bench/lux.jl
+++ b/bench/lux.jl
@@ -34,14 +34,14 @@ for i in 1:5
     chs = (1, 128, fill(64, i)..., 128, 1)
     model = FourierNeuralOperator(gelu; chs=chs, modes=(16,))
     ps, st = Lux.setup(rng, model)
-    _ = model(x, ps, st) # TTFX
+    model(x, ps, st) # TTFX
 
     t_fwd[i] = @belapsed $model($x, $ps, $st)
 
     t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
 end
 
-println("\n ### FNO")
+println("\n## FNO")
 print("| #layers | Forward | Train: 10 epochs | \n")
 print("| --- | --- | --- | \n")
 for i in 1:5
@@ -67,18 +67,16 @@ for i in 1:5
     ch_trunk = (dim_y, fill(64, i)..., 128)
     model = DeepONet(; branch=ch_branch, trunk=ch_trunk)
     ps, st = Lux.setup(rng, model)
-    _ = model((u, y), ps, st) # TTFX
+    model((u, y), ps, st) # TTFX
 
     t_fwd[i] = @belapsed $model(($u, $y), $ps, $st)
 
     t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
 end
 
-println("\n ### DeepONet")
+println("\n## DeepONet")
 print("| #layers | Forward | Train: 10 epochs | \n")
 print("| --- | --- | --- | \n")
 for i in 1:5
     print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
 end
-
-