Powered by AppSignal & Oban Pro
Would you like to see your link here? Contact us

Stop Reinventing the Wheel

StopReinventingTheWheel.livemd

Stop Reinventing the Wheel

Mix.install([
  {:axon_onnx, "~> 0.4"},
  {:axon, "~> 0.5"},
  {:nx, "~> 0.5"},
  {:exla, "~> 0.5"},
  {:stb_image, "~> 0.6"},
  {:kino, "~> 0.8"}
])
Resolving Hex dependencies...
Resolution completed in 0.242s
New:
  axon 0.5.1
  axon_onnx 0.4.0
  cc_precompiler 0.1.6
  complex 0.5.0
  decimal 2.0.0
  elixir_make 0.7.5
  exla 0.5.1
  kino 0.8.1
  nx 0.5.1
  protox 1.6.10
  stb_image 0.6.0
  table 0.1.2
  telemetry 1.2.1
  xla 0.4.4
* Getting axon_onnx (Hex package)
* Getting axon (Hex package)
* Getting nx (Hex package)
* Getting exla (Hex package)
* Getting stb_image (Hex package)
* Getting kino (Hex package)
* Getting table (Hex package)
* Getting cc_precompiler (Hex package)
* Getting elixir_make (Hex package)
* Getting telemetry (Hex package)
* Getting xla (Hex package)
* Getting complex (Hex package)
* Getting protox (Hex package)
* Getting decimal (Hex package)
==> decimal
Compiling 4 files (.ex)
Generated decimal app
==> table
Compiling 5 files (.ex)
Generated table app
===> Analyzing applications...
===> Compiling telemetry
==> protox
Compiling 53 files (.ex)
Generated protox app
==> complex
Compiling 2 files (.ex)
Generated complex app
==> nx
Compiling 31 files (.ex)
Generated nx app
==> kino
Compiling 37 files (.ex)
Generated kino app
==> axon
Compiling 23 files (.ex)
Generated axon app
==> axon_onnx
Compiling 28 files (.ex)
warning: Nx.power/2 is deprecated. Use pow/2 instead
  lib/axon_onnx/deserialize.ex:364: AxonOnnx.Deserialize

warning: Nx.power/2 is deprecated. Use pow/2 instead
Invalid call found at 4 locations:
  lib/axon_onnx/shared.ex:44: AxonOnnx.Shared."__defn:sumsquare__"/2
  lib/axon_onnx/shared.ex:52: AxonOnnx.Shared."__defn:l2_norm__"/2
  lib/axon_onnx/shared.ex:63: AxonOnnx.Shared."__defn:lrn__"/2
  lib/axon_onnx/shared.ex:65: AxonOnnx.Shared."__defn:lrn__"/2

warning: Nx.Defn.Kernel.transform/2 is deprecated. Use deftransform/2 or deftransformp/2 from Nx.Defn instead
Invalid call found at 2 locations:
  lib/axon_onnx/shared.ex:58: AxonOnnx.Shared."__defn:lrn__"/2
  lib/axon_onnx/shared.ex:94: AxonOnnx.Shared."__defn:numpy_matmul__"/3

warning: Nx.power/2 is deprecated. Use pow/2 instead
  lib/axon_onnx/deserialize.ex: AxonOnnx.Deserialize.recur_nodes/2

warning: Nx.random_normal/4 is deprecated. Use Nx.Random instead
Invalid call found at 4 locations:
  lib/axon_onnx/deserialize.ex:1978: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:2009: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:2015: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:2022: AxonOnnx.Deserialize.recur_nodes/2

warning: Nx.random_uniform/4 is deprecated. Use Nx.Random.uniform/2 instead
Invalid call found at 4 locations:
  lib/axon_onnx/deserialize.ex:1911: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:1942: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:1948: AxonOnnx.Deserialize.recur_nodes/2
  lib/axon_onnx/deserialize.ex:1955: AxonOnnx.Deserialize.recur_nodes/2

Generated axon_onnx app
==> elixir_make
Compiling 6 files (.ex)
Generated elixir_make app
==> cc_precompiler
Compiling 3 files (.ex)
Generated cc_precompiler app
==> stb_image
Compiling 2 files (.ex)
Generated stb_image app
==> xla
Compiling 2 files (.ex)
Generated xla app
==> exla
Unpacking /Users/sean/Library/Caches/xla/0.4.4/cache/download/xla_extension-aarch64-darwin-cpu.tar.gz into /Users/sean/Library/Caches/mix/installs/elixir-1.14.2-erts-13.0.2/766a4043250769d01a3a59ff6d42c870/deps/exla/cache
Using libexla.so from /Users/sean/Library/Caches/xla/exla/elixir-1.14.2-erts-13.0.2-xla-0.4.4-exla-0.5.1-44u5ai52odx7l3gshk4r67yj4y/libexla.so
Compiling 21 files (.ex)
Generated exla app
:ok

Input Data

defmodule CatsAndDogs do
  def pipeline(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&to_tensors(&1, target_height, target_width))
    |> Stream.chunk_every(batch_size)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  def pipeline_with_augmentations(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&to_tensors(&1, target_height, target_width))
    |> Stream.map(&random_flip(&1, :height))
    |> Stream.map(&random_flip(&1, :width))
    |> Stream.chunk_every(batch_size)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  defp random_flip({image, label}, axis) do
    if :rand.uniform() < 0.5 do
      {Nx.reverse(image, axes: [axis]), label}
    else
      {image, label}
    end
  end

  defp parse_image(path) do
    label = if String.contains?(path, "cat"), do: 0, else: 1

    case StbImage.read_file(path) do
      {:ok, img} -> {img, label}
      _error -> :error
    end
  end

  defp to_tensors({:ok, {img, label}}, target_height, target_width) do
    img_tensor =
      img
      |> StbImage.resize(target_height, target_width)
      |> StbImage.to_nx()
      |> Nx.divide(255)
      |> Nx.transpose(axes: [:channels, :height, :width])

    label_tensor = Nx.tensor([label])

    {img_tensor, label_tensor}
  end
end
{:module, CatsAndDogs, <<70, 79, 82, 49, 0, 0, 21, ...>>, {:to_tensors, 3}}
{test_paths, train_paths} =
  Path.wildcard("train/*.jpg")
  |> Enum.shuffle()
  |> Enum.split(1000)

{test_paths, val_paths} = test_paths |> Enum.split(750)

batch_size = 32
target_height = 160
target_width = 160

train_pipeline =
  CatsAndDogs.pipeline_with_augmentations(
    train_paths,
    batch_size,
    target_height,
    target_width
  )

val_pipeline =
  CatsAndDogs.pipeline(
    val_paths,
    batch_size,
    target_height,
    target_width
  )

test_pipeline =
  CatsAndDogs.pipeline(
    test_paths,
    batch_size,
    target_height,
    target_width
  )

Enum.take(train_pipeline, 1)
[
  {#Nx.Tensor<
     f32[32][channels: 3][height: 160][width: 160]
     [
       [
         [
           [0.250980406999588, 0.2549019753932953, 0.26274511218070984, 0.2666666805744171, 0.2705882489681244, 0.27843138575553894, 0.29019609093666077, 0.2980392277240753, 0.30980393290519714, 0.30588236451148987, 0.3137255012989044, 0.3294117748737335, 0.34117648005485535, 0.3490196168422699, 0.3490196168422699, 0.3529411852359772, 0.35686275362968445, 0.3607843220233917, 0.3607843220233917, 0.3607843220233917, 0.3529411852359772, 0.3607843220233917, 0.3803921639919281, 0.38823530077934265, 0.3843137323856354, 0.38823530077934265, 0.3921568691730499, 0.3921568691730499, 0.3960784375667572, 0.40392157435417175, 0.40784314274787903, 0.4117647111415863, 0.4117647111415863, 0.4117647111415863, 0.4117647111415863, 0.41960784792900085, 0.42352941632270813, 0.42352941632270813, 0.42352941632270813, 0.4274509847164154, 0.43529412150382996, 0.43921568989753723, 0.43529412150382996, 0.43921568989753723, 0.4470588266849518, 0.45490196347236633, 0.45490196347236633, 0.45490196347236633, ...],
           ...
         ],
         ...
       ],
       ...
     ]
   >,
   #Nx.Tensor<
     s64[32][1]
     [
       [0],
       [0],
       [0],
       [1],
       [0],
       [1],
       [0],
       [0],
       [1],
       [1],
       [0],
       [0],
       [0],
       [0],
       [0],
       [1],
       [1],
       [0],
       [0],
       [1],
       [0],
       [1],
       [0],
       [0],
       [0],
       [1],
       [1],
       [1],
       [0],
       [0],
       [0],
       [1]
     ]
   >}
]

Feature Extraction

{cnn_base, cnn_base_params} = AxonOnnx.import("mobilenetv2-7.onnx", batch_size: batch_size)
{#Axon<
   inputs: %{"data" => {1, 3, 224, 224}}
   outputs: "mobilenetv20_output_flatten0_reshape0"
   nodes: 166
 >,
 %{
   "mobilenetv20_features_linearbottleneck14_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[960]
       [-0.052753571420907974, -0.09027774631977081, -0.10990642756223679, 0.024725452065467834, -0.019158951938152313, 0.03750738501548767, -0.0568571463227272, 0.055139604955911636, -0.0301227867603302, 0.01030650082975626, 0.14296390116214752, 0.030574854463338852, -0.0474417544901371, -0.01992128975689411, -0.011829576455056667, 0.03172306343913078, 0.026741042733192444, -0.055858079344034195, 0.02560657635331154, -0.05073458328843117, -0.01601296477019787, -0.031947072595357895, 0.2531116306781769, -0.11873234063386917, -0.15965522825717926, -0.014474667608737946, 0.13942478597164154, 0.054266903549432755, 0.09119295328855515, -0.08511218428611755, 0.06783539801836014, -0.18136344850063324, -0.054311662912368774, -0.054432209581136703, -0.032984793186187744, 0.03091847524046898, -0.009745757095515728, -0.052780959755182266, -0.05558377131819725, -0.05194559320807457, 0.0438227616250515, -0.006544236559420824, -0.010287697426974773, 0.24488094449043274, -0.014213151298463345, 0.03649181127548218, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[960]
       [0.20604729652404785, 0.22227846086025238, 0.19922102987766266, 0.15634967386722565, 0.17543338239192963, 0.15423689782619476, 0.209830641746521, 0.16449765861034393, 0.18006004393100739, 0.16814470291137695, 0.128125861287117, 0.12089566886425018, 0.21123994886875153, 0.15545915067195892, 0.1902669370174408, 0.16879144310951233, 0.18586121499538422, 0.20589303970336914, 0.1682288497686386, 0.17667225003242493, 0.1531776338815689, 0.1845463514328003, 0.09342513978481293, 0.23671939969062805, 0.192972794175148, 0.1480586677789688, 0.13556815683841705, 0.1454746574163437, 0.12987960875034332, 0.2190699428319931, 0.14040623605251312, 0.4639834761619568, 0.26004794239997864, 0.22225748002529144, 0.18369433283805847, 0.16256587207317352, 0.1834648698568344, 0.21532562375068665, 0.23163369297981262, 0.23751749098300934, 0.13892124593257904, 0.17630253732204437, 0.17822310328483582, 0.15526926517486572, 0.20411202311515808, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[960]
       [-0.0023206884507089853, -0.004551870748400688, 0.0034389428328722715, 0.0015984668862074614, -0.002293417463079095, -0.03154199197888374, 0.015855349600315094, -0.021844547241926193, -0.0030182870104908943, -0.008290023542940617, -0.008337621577084064, 0.003635477740317583, -0.02967023104429245, 0.009609879925847054, -0.00535233598202467, -0.03133776783943176, -0.034155990928411484, -0.0017219285946339369, -0.02915215492248535, -0.0034448853693902493, -0.0022068622056394815, -0.003281062701717019, 0.003420598339289427, 0.00949984323233366, 8.879394154064357e-4, 0.002022636355832219, -0.02363204024732113, -0.017934758216142654, -0.02237909846007824, -0.05140375345945358, 4.167649312876165e-4, -0.006942620035260916, 0.011951959691941738, -0.0031885183416306973, -9.77475312538445e-4, -0.0014110811753198504, -0.03238283097743988, -0.005639346782118082, -0.0070093655958771706, -0.005301866214722395, 0.0020262207835912704, -0.03578586131334305, -0.0526641346514225, -0.009258503094315529, ...]
     >,
     "var" => #Nx.Tensor<
       f32[960]
       [4.962424281984568e-4, 2.4244154337793589e-4, 2.3794326989445835e-4, 2.723340585362166e-4, 3.9763047243468463e-4, 5.111433565616608e-4, 4.7444249503314495e-4, 3.682477108668536e-4, 4.1800420149229467e-4, 7.685577729716897e-4, 3.6789028672501445e-4, 2.2397300926968455e-4, 4.4611067278310657e-4, 2.4524933542124927e-4, 0.0011620114091783762, 5.935822846367955e-4, 7.234420045278966e-4, 3.640245122369379e-4, 5.758852930739522e-4, 3.4469107049517334e-4, 2.2663685376755893e-4, 4.212635103613138e-4, 2.179306757170707e-4, 4.0046244976110756e-4, 1.659863191889599e-4, 8.642849934403785e-6, 7.320880540646613e-4, 3.3074370003305376e-4, 4.445896774996072e-4, 7.599603850394487e-4, 3.4124808735214174e-4, 0.0010172502370551229, 7.496927282772958e-4, 4.7150562750175595e-4, 7.05038895830512e-4, 2.509213227313012e-4, 5.495413788594306e-4, 4.112518217880279e-4, 0.0011761350324377418, 6.135390140116215e-4, 4.0227093268185854e-4, 6.605967064388096e-4, 5.76652237214148e-4, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck15_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[160][960][1][1]
       [
         [
           [
             [0.0035520142409950495]
           ],
           [
             [-0.09606171399354935]
           ],
           [
             [0.0023302228655666113]
           ],
           [
             [0.1408250778913498]
           ],
           [
             [0.029948130249977112]
           ],
           [
             [-0.05268537625670433]
           ],
           [
             [-0.03461449220776558]
           ],
           [
             [-0.02320627123117447]
           ],
           [
             [0.046411626040935516]
           ],
           [
             [-0.010943359695374966]
           ],
           [
             [0.02103927917778492]
           ],
           [
             [-0.04870555177330971]
           ],
           [
             [0.01763751357793808]
           ],
           [
             [0.02220737934112549]
           ],
           [
             [-0.012389437295496464]
           ],
           [
             [0.024784402921795845]
           ],
           [
             [-0.023123491555452347]
           ],
           [
             [-0.032055873423814774]
           ],
           [
             [-0.0038467550184577703]
           ],
           [
             [-0.027905654162168503]
           ],
           [
             [0.011424683034420013]
           ],
           [
             [0.008500807918608189]
           ],
           [
             [-0.053142279386520386]
           ],
           [
             [-0.0061360690742731094]
           ],
           [
             [-0.032072391360998154]
           ],
           [
             [-0.01983955316245556]
           ],
           [
             [0.12456362694501877]
           ],
           [
             [0.05674683302640915]
           ],
           [
             [-0.012410887517035007]
           ],
           [
             [-0.037049565464258194]
           ],
           [
             [-0.05917417258024216]
           ],
           [
             [0.03232760727405548]
           ],
           [
             [-0.017962725833058357]
           ],
           [
             [-0.04565010964870453]
           ],
           [
             [-0.0036436335649341345]
           ],
           [
             [-0.054335057735443115]
           ],
           [
             [-0.01371700968593359]
           ],
           [
             [-0.0356641449034214]
           ],
           [
             [-0.07494678348302841]
           ],
           [
             [-0.027322199195623398]
           ],
           [
             [-0.04358804598450661]
           ],
           [
             [0.04669416323304176]
           ],
           [
             [0.028277995064854622]
           ],
           [
             [-0.04000042378902435]
           ],
           [
             [0.005873105954378843]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck12_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][576][1][1]
       [
         [
           [
             [0.03595563769340515]
           ],
           [
             [0.0247351061552763]
           ],
           [
             [0.049286190420389175]
           ],
           [
             [0.12805162370204926]
           ],
           [
             [-0.0414484404027462]
           ],
           [
             [-0.06513191014528275]
           ],
           [
             [-0.03178936615586281]
           ],
           [
             [-0.06259823590517044]
           ],
           [
             [0.10336058586835861]
           ],
           [
             [0.045249905437231064]
           ],
           [
             [0.05266702175140381]
           ],
           [
             [0.026978133246302605]
           ],
           [
             [-0.007294573821127415]
           ],
           [
             [-0.0391872338950634]
           ],
           [
             [0.05743891000747681]
           ],
           [
             [-0.0659816786646843]
           ],
           [
             [0.005016516428440809]
           ],
           [
             [-0.006897100247442722]
           ],
           [
             [-0.05635765939950943]
           ],
           [
             [0.02955286018550396]
           ],
           [
             [0.0061321984976530075]
           ],
           [
             [-0.07741623371839523]
           ],
           [
             [-0.07541409134864807]
           ],
           [
             [0.03161944821476936]
           ],
           [
             [-0.06556887179613113]
           ],
           [
             [-0.036827001720666885]
           ],
           [
             [-0.04910659417510033]
           ],
           [
             [0.011762046255171299]
           ],
           [
             [-0.008211801759898663]
           ],
           [
             [0.030010191723704338]
           ],
           [
             [0.020299093797802925]
           ],
           [
             [0.07257092744112015]
           ],
           [
             [0.01673484407365322]
           ],
           [
             [0.07540816813707352]
           ],
           [
             [-0.04706030339002609]
           ],
           [
             [-0.0737130269408226]
           ],
           [
             [0.01491565816104412]
           ],
           [
             [-0.05089497193694115]
           ],
           [
             [-0.012679836712777615]
           ],
           [
             [0.08336318284273148]
           ],
           [
             [0.0022198110818862915]
           ],
           [
             [-0.04258105531334877]
           ],
           [
             [-0.0074499910697340965]
           ],
           [
             [0.048486143350601196]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck14_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[960]
       [0.11810629814863205, 0.11761260032653809, -0.026628313586115837, 0.06580127030611038, 0.0590800978243351, 0.0816112533211708, -0.004318615887314081, 0.06424110382795334, 0.07805308699607849, 0.014792101457715034, 0.03583914414048195, -0.04613340273499489, 0.09175406396389008, -0.016964279115200043, 0.12577180564403534, 0.0832546055316925, 0.0908178836107254, 0.09375405311584473, 0.10626447200775146, 0.005106383468955755, 0.016398748382925987, 0.07167958468198776, -0.038985028862953186, 0.015483750030398369, 0.018431706354022026, -0.012954303063452244, 0.013924377039074898, 0.02286217361688614, 0.038238342851400375, 0.1492835134267807, 0.02409234456717968, -0.055206529796123505, -0.06437041610479355, 0.1014975979924202, 0.09384036064147949, 0.1150284856557846, 0.09201668202877045, 0.01461042370647192, 0.11383866518735886, 0.1275649517774582, 0.07282229512929916, 0.10276996344327927, 0.1410975307226181, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[960]
       [0.08687656372785568, 0.061947379261255264, 0.11407668143510818, 0.10592024773359299, 0.12367846071720123, 0.11453951895236969, 0.1292855590581894, 0.12500832974910736, 0.09852098673582077, 0.1363425850868225, 0.11740787327289581, 0.12203019112348557, 0.11298535764217377, 0.1157105565071106, 0.11476130038499832, 0.12099182605743408, 0.13221073150634766, 0.08549683541059494, 0.12948764860630035, 0.12371236085891724, 0.11909203976392746, 0.10468459874391556, 0.12764997780323029, 0.10788445919752121, 0.1629999279975891, 0.038849350064992905, 0.17148640751838684, 0.14035196602344513, 0.1376633644104004, 0.10015642642974854, 0.12715141475200653, 0.17000487446784973, 0.16757212579250336, 0.08869356662034988, 0.11598992347717285, 0.07029741257429123, 0.13028739392757416, 0.12240685522556305, 0.10917554795742035, 0.07950529456138611, 0.11374188214540482, 0.12892962992191315, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[960]
       [-1.49989478614998e-8, -6.598725121875759e-8, -3.3777659780298563e-8, 6.155921328243608e-10, 1.6680649084221955e-9, -6.975810151743644e-8, -1.0280804474405159e-7, 1.1829909851712728e-7, -1.391223776181505e-7, 1.951114825260447e-8, -3.154044847519799e-8, -7.882130859115932e-8, 7.139960445101678e-8, -6.877448299746902e-8, -1.5163976740950602e-7, 4.35149836164328e-8, -3.990682628796094e-8, 3.889567778969649e-8, 8.67010143679181e-8, 1.5913511219878274e-7, -1.216080569577116e-8, -3.217380850628615e-8, -2.502063409792754e-7, -1.2450051656287542e-8, -8.94942004947552e-8, -1.0849137055402025e-7, -7.23792865642281e-8, -2.6819417797696588e-9, -4.9953573721950306e-8, -8.798366479823017e-8, 4.691241528576029e-8, 1.7972880073102715e-8, -1.1659428622579071e-7, -1.6043891548633837e-8, 1.5714466883309797e-7, 1.3846912061410421e-8, -8.361175218851713e-8, -2.2993779325020114e-8, -2.5269557557550115e-9, 2.5268136027989385e-8, -5.639618905206589e-8, ...]
     >,
     "var" => #Nx.Tensor<
       f32[960]
       [0.015375461429357529, 0.008966661989688873, 0.014436393044888973, 0.018274905160069466, 0.022764142602682114, 0.017717313021421432, 0.028602011501789093, 0.02376588247716427, 0.013949170708656311, 0.014504922553896904, 0.015037551522254944, 0.017101136967539787, 0.014324591495096684, 0.024719757959246635, 0.022843269631266594, 0.016248025000095367, 0.02265304885804653, 0.013783901929855347, 0.020945660769939423, 0.010627896524965763, 0.02857845090329647, 0.019863925874233246, 0.010918163694441319, 0.01824241690337658, 0.061989493668079376, 0.013930430635809898, 0.03650098294019699, 0.01744050718843937, 0.015627671033143997, 0.011513426899909973, 0.020515792071819305, 0.024865739047527313, 0.028932511806488037, 0.013199539855122566, 0.01918703131377697, 0.008621279150247574, 0.022145573049783707, 0.017280828207731247, 0.022399913519620895, 0.01449136808514595, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck16_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[960]
       [0.04987475648522377, 0.07547232508659363, 0.05202540010213852, 0.3060879707336426, 0.061899781227111816, 0.19728609919548035, 0.05910351499915123, 0.05794937163591385, 0.08533387631177902, 0.059474583715200424, 0.07320644706487656, 0.20866796374320984, 0.0625239685177803, 0.011463859118521214, 0.14914481341838837, 0.05909585580229759, 0.05309757962822914, 0.06794954091310501, 0.059355489909648895, 0.0552368201315403, 0.2624534070491791, 0.06431957334280014, 0.06486242264509201, 0.07900767773389816, 0.06663036346435547, 0.0600312314927578, 0.16903376579284668, 0.11807416379451752, 0.061896178871393204, 0.06435873359441757, 0.06086957827210426, 0.0652950331568718, 0.058827582746744156, -0.12069928646087646, 0.0629323199391365, 0.07213398069143295, 0.061805903911590576, 0.06384092569351196, 0.12417661398649216, 0.04771304503083229, 0.06310302764177322, 0.05863166227936745, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[960]
       [0.23864960670471191, 0.2322680801153183, 0.2386099249124527, 0.09403114020824432, 0.22456052899360657, 0.18399876356124878, 0.2542182505130768, 0.2437894642353058, 0.1435527354478836, 0.2224864661693573, 0.2270522117614746, 0.14962053298950195, 0.26540496945381165, 0.2219705581665039, 0.1094207614660263, 0.18297213315963745, 0.27974048256874084, 0.23882733285427094, 0.2643541991710663, 0.23397229611873627, 0.09304135292768478, 0.2347753345966339, 0.15721535682678223, 0.14321495592594147, 0.21017815172672272, 0.2423482984304428, 0.13239087164402008, 0.1255304217338562, 0.21942436695098877, 0.21238590776920319, 0.19641195237636566, 0.18220824003219604, 0.18533524870872498, 0.2664465010166168, 0.22679947316646576, 0.22101089358329773, 0.2273634821176529, 0.21875879168510437, 0.12264074385166168, 0.13991476595401764, 0.22431768476963043, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[960]
       [3.730551979970187e-4, 5.662986077368259e-4, 3.1321379356086254e-4, -0.006701103411614895, 5.611852975562215e-4, -0.0017290484393015504, 3.5307201324030757e-4, 3.786919405683875e-4, -0.010407516732811928, 4.5713825966231525e-4, 5.702645285055041e-4, -0.0018510040827095509, 3.712004399858415e-4, -0.05779576674103737, -0.012207098305225372, 7.817225414328277e-4, 2.1997511794324964e-4, 4.943913081660867e-4, 2.8153584571555257e-4, 3.6745445686392486e-4, -0.014429425820708275, 4.293794627301395e-4, -0.007664929144084454, -0.010241005569696426, 6.079634767957032e-4, 4.3814798118546605e-4, -0.008106333203613758, -0.007122342940419912, 5.337817710824311e-4, 5.633817054331303e-4, 7.095154142007232e-4, 7.763216271996498e-4, 7.150981109589338e-4, -0.007433908525854349, 4.647991736419499e-4, 6.517147412523627e-4, 4.5494179357774556e-4, 5.59471664018929e-4, -0.01687384955585003, -0.01668229140341282, ...]
     >,
     "var" => #Nx.Tensor<
       f32[960]
       [1.1024731065845117e-6, 1.4655657878392958e-6, 9.686918929219246e-7, 3.663915776996873e-5, 1.61286152433604e-6, 6.336554179142695e-6, 9.814286840992281e-7, 1.0332710189686622e-6, 1.4501674741040915e-4, 1.495523747507832e-6, 1.517296595920925e-6, 6.068381480872631e-6, 1.2031182450300548e-6, 7.863730425015092e-4, 1.2193145812489092e-4, 2.1900182218814734e-6, 6.346336363094451e-7, 1.3137222367731738e-6, 8.681112717567885e-7, 1.2658803143494879e-6, 1.3425096403807402e-4, 1.1334504961268976e-6, 6.453898822655901e-5, 1.1336532043060288e-4, 1.7878132894111332e-6, 1.2046853044012096e-6, 8.614131365902722e-5, 1.0490042041055858e-4, 1.5527492678302224e-6, 1.5868641867200495e-6, 1.6290416624542559e-6, 2.0585173388099065e-6, 1.81840641744202e-6, 2.3889367002993822e-4, 1.2654711554205278e-6, 1.7507608163214172e-6, 1.2430069773472496e-6, 1.4594443200621754e-6, 2.228585071861744e-4, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck4_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[192]
       [-0.16025467216968536, -0.08318334817886353, -0.013605363667011261, 0.016528524458408356, -0.0414351262152195, -0.08624469488859177, 0.08369371294975281, 0.08077652007341385, -0.3613238036632538, 0.0863688737154007, -0.07133538275957108, 0.17072106897830963, -0.006156732328236103, 0.04610714688897133, 0.26338741183280945, 0.25080469250679016, 0.0866481363773346, 0.029331553727388382, 0.13819074630737305, -0.06575538963079453, -5.868459120392799e-4, -0.02378389984369278, -0.03555862605571747, -0.022012412548065186, 0.07437903434038162, 0.07370810955762863, 0.30864641070365906, -0.006008080206811428, 0.15304052829742432, -0.00992464367300272, 0.19627465307712555, 0.07678351551294327, 0.10757749527692795, -0.02487078309059143, -0.040003497153520584, -0.07770790904760361, 0.2532285451889038, 0.16046945750713348, 0.05568472295999527, 0.15549562871456146, -0.038443390280008316, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[192]
       [0.12059392035007477, 0.18736249208450317, 0.1590438038110733, 0.2091878354549408, 0.1450766772031784, 0.15702469646930695, 0.3902188539505005, 0.25355497002601624, 0.3910222351551056, 0.26179036498069763, 0.189780130982399, 0.14822576940059662, 0.210128515958786, 0.1974102109670639, 0.0950021743774414, 0.1603325754404068, 0.12979912757873535, 0.08638784289360046, 0.11244207620620728, 0.1680508255958557, 0.243540421128273, 0.20751050114631653, 0.1465093046426773, 0.15335243940353394, 0.2690037190914154, 0.1203547790646553, 0.0906066969037056, 0.23154661059379578, 0.1154584214091301, 0.1968846470117569, 0.11754576861858368, 0.1189347431063652, 0.3180309534072876, 0.20275117456912994, 0.1695779711008072, 0.2201462984085083, 0.10936479270458221, 0.0758633241057396, 0.12570320069789886, 0.09598797559738159, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[192]
       [0.005834498442709446, 0.005320976488292217, 0.011644497513771057, -0.0033567524515092373, 0.006512106861919165, 0.0075160483829677105, -0.011912024579942226, -0.021598124876618385, -0.15417402982711792, -0.001387530704960227, -0.0037699746899306774, 0.01811889372766018, -9.935529669746757e-4, -0.00730853620916605, -0.004980568308383226, 5.785240791738033e-4, 0.0198814794421196, 0.0036787642166018486, -0.018057825043797493, -2.4293076421599835e-4, 0.002015749691054225, 0.013546604663133621, 0.007414540275931358, 1.2083808542229235e-4, -0.003124411916360259, -0.019600532948970795, -0.005146954674273729, -0.002622810425236821, -0.016293222084641457, 0.023455090820789337, -0.00871172547340393, -0.008095859549939632, -0.018160350620746613, -5.122238071635365e-4, 0.0051345909014344215, 0.002093162387609482, -0.01451784186065197, 0.0019358042627573013, 0.0016494555165991187, ...]
     >,
     "var" => #Nx.Tensor<
       f32[192]
       [1.82411604328081e-4, 2.6851065922528505e-4, 2.388732973486185e-4, 8.272912818938494e-4, 1.7400244541931897e-4, 2.264920767629519e-4, 0.005600188858807087, 0.0019779913127422333, 0.004019826650619507, 0.002996259368956089, 2.9179989360272884e-4, 9.683401440270245e-4, 6.0461163229774684e-5, 1.024205848807469e-4, 3.014750254806131e-4, 4.56174939245102e-6, 8.186800550902262e-5, 4.768615690409206e-5, 3.702163230627775e-4, 3.5265996120870113e-4, 0.001435167738236487, 4.292532103136182e-4, 1.5066159539856017e-4, 1.9934076408389956e-4, 0.0030081914737820625, 9.441975998925045e-5, 3.100578032899648e-4, 2.9591855127364397e-4, 4.2446120642125607e-4, 8.81795072928071e-4, 1.9674932991620153e-4, 2.7562043396756053e-4, 0.0034803457092493773, 6.995198782533407e-4, 3.369579208083451e-4, 6.203856319189072e-4, 4.934457829222083e-4, 4.415657167555764e-5, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck13_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[576][1][3][3]
       [
         [
           [
             [0.21406474709510803, 0.12984389066696167, -0.1617448329925537],
             [0.19517506659030914, -0.11259430646896362, 0.024987775832414627],
             [-0.1077454686164856, 0.06464696675539017, -0.7981051802635193]
           ]
         ],
         [
           [
             [-1.1268330812454224, 0.10583485662937164, 0.04849735647439957],
             [-0.1629934459924698, 0.1864052414894104, -0.374404639005661],
             [0.18732739984989166, 0.03825368359684944, 0.3030540645122528]
           ]
         ],
         [
           [
             [-0.1971191167831421, 0.21261656284332275, 0.01313933078199625],
             [0.15392960608005524, 0.09030754119157791, 0.06439008563756943],
             [-0.7180318236351013, -0.04540767893195152, 0.44523710012435913]
           ]
         ],
         [
           [
             [-0.12377922981977463, 0.16086746752262115, 0.35277417302131653],
             [0.5762680172920227, 0.037068989127874374, 0.523024320602417],
             [0.2114313840866089, 0.6446085572242737, 0.36676862835884094]
           ]
         ],
         [
           [
             [0.6760067939758301, -0.05647632107138634, 0.04974052682518959],
             [0.060135506093502045, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck1_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[24][96][1][1]
       [
         [
           [
             [0.09949389100074768]
           ],
           [
             [-0.18654611706733704]
           ],
           [
             [0.13225184381008148]
           ],
           [
             [0.03560471162199974]
           ],
           [
             [-0.020650217309594154]
           ],
           [
             [-0.06614962220191956]
           ],
           [
             [-0.12619785964488983]
           ],
           [
             [-0.17344357073307037]
           ],
           [
             [0.13648797571659088]
           ],
           [
             [0.05030941590666771]
           ],
           [
             [0.0744558572769165]
           ],
           [
             [0.02734171226620674]
           ],
           [
             [0.1788010597229004]
           ],
           [
             [-0.0010324945906177163]
           ],
           [
             [-0.058398596942424774]
           ],
           [
             [-0.30487778782844543]
           ],
           [
             [0.23805657029151917]
           ],
           [
             [-0.04994773864746094]
           ],
           [
             [-0.04996243491768837]
           ],
           [
             [0.049236271530389786]
           ],
           [
             [0.15690600872039795]
           ],
           [
             [-0.085136778652668]
           ],
           [
             [-0.06452983617782593]
           ],
           [
             [-0.03775155171751976]
           ],
           [
             [-0.2167663425207138]
           ],
           [
             [0.20706818997859955]
           ],
           [
             [0.23726293444633484]
           ],
           [
             [-0.1283334344625473]
           ],
           [
             [-0.22454236447811127]
           ],
           [
             [-0.18676628172397614]
           ],
           [
             [0.09907364100217819]
           ],
           [
             [0.22619308531284332]
           ],
           [
             [-0.029751716181635857]
           ],
           [
             [0.16340412199497223]
           ],
           [
             [0.01070564053952694]
           ],
           [
             [0.3250269889831543]
           ],
           [
             [0.21018050611019135]
           ],
           [
             [-0.062177326530218124]
           ],
           [
             [0.12571603059768677]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck6_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[192]
       [0.3632318675518036, -0.023602070286870003, -0.0614071749150753, 0.5761924386024475, 0.1690974086523056, 0.19036972522735596, 0.0031085966620594263, -0.018681762740015984, 0.038269754499197006, 0.03641441464424133, 0.029138026759028435, -0.016922323033213615, 0.055098824203014374, 0.21938563883304596, 0.22719338536262512, 0.20525230467319489, 0.20284974575042725, 0.07362302392721176, 0.1064985990524292, 0.40782883763313293, -0.00754354614764452, 0.009455508552491665, 0.4253721535205841, 0.6265411972999573, 0.22528603672981262, -0.05191487818956375, 0.6341891288757324, 0.347513884305954, 0.04142335429787636, 0.08137086033821106, 0.35836175084114075, 0.038170408457517624, 0.14619027078151703, 0.0850687101483345, 0.36599797010421753, 0.3969757854938507, 0.09776335954666138, 0.02894003875553608, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[192]
       [0.1283191293478012, 0.2977066934108734, 0.2804311215877533, 0.20587405562400818, 0.2696685492992401, 0.12131325900554657, 0.30051329731941223, 0.36318743228912354, 0.23120905458927155, 0.22658991813659668, 0.24182915687561035, 0.30583131313323975, 0.20871473848819733, 0.2518334686756134, 0.22621870040893555, 0.20286978781223297, 0.19829389452934265, 0.16485567390918732, 0.18886375427246094, 0.1450076848268509, 0.312345415353775, 0.2126840353012085, 0.2775242030620575, 0.1377149373292923, 0.15688155591487885, 0.3647836446762085, 0.24961578845977783, 0.18082654476165771, 0.17116306722164154, 0.23007430136203766, 0.1871042400598526, 0.2649954557418823, 0.306458443403244, 0.21913877129554749, 0.16099828481674194, 0.1434336155653, 0.12895677983760834, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[192]
       [-0.04431454464793205, 0.017416827380657196, 0.006211827043443918, 0.06156991794705391, -0.1327718049287796, -0.038969025015830994, -0.16951747238636017, 0.07648631185293198, -0.09081843495368958, 0.0041724517941474915, 0.02877575159072876, -0.015150758437812328, 0.03303484991192818, -0.07873229682445526, -0.06547018140554428, 0.0404805913567543, 0.04381240904331207, 0.023880552500486374, -0.05104473978281021, -0.035787083208560944, 0.03322664275765419, 0.022716958075761795, 0.016515308991074562, -0.0348416231572628, -0.04564644396305084, -0.027626102790236473, 0.09955187886953354, -0.06939177215099335, 0.05254083871841431, 0.04370502755045891, 0.034599799662828445, -0.14765559136867523, 0.023225821554660797, 0.05121980607509613, 0.0666728988289833, -0.010222886689007282, ...]
     >,
     "var" => #Nx.Tensor<
       f32[192]
       [0.001848297892138362, 0.0011715174186974764, 0.0011900259414687753, 0.005860364064574242, 0.012526428326964378, 0.0012380351545289159, 0.0028932117857038975, 9.825429879128933e-4, 5.667941877618432e-4, 0.002237679436802864, 0.001042675692588091, 0.0024795986246317625, 0.0012960642343387008, 0.01065833680331707, 0.0063052731566131115, 0.008763182908296585, 0.0048617287538945675, 0.001241417950950563, 0.0015297426143661141, 0.002018269617110491, 0.0011946061858907342, 0.0011670668609440327, 0.001651656930334866, 0.001621170318685472, 0.001029442879371345, 9.477304411120713e-4, 0.01142901461571455, 7.548464927822351e-4, 0.0023103056009858847, 0.0028397291898727417, 0.0016283183358609676, 0.004941418766975403, 0.0020476358477026224, 0.002869980875402689, 0.0012202837970107794, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck1_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[24]
       [-5.599617361440323e-5, -1.388202654197812e-4, 1.4964351430535316e-4, -4.48540726210922e-4, -2.7302927264827304e-5, 1.428744290024042e-4, -1.0991278395522386e-4, 1.3559161743614823e-4, 1.4486558211501688e-4, -5.513934302143753e-4, -9.67518353718333e-5, -2.075896190945059e-4, -1.0409099922981113e-4, -1.4599738642573357e-4, -1.6489351401105523e-4, -9.032035450218245e-5, -1.591532345628366e-4, -1.570975873619318e-4, -7.89462574175559e-5, 3.225041655241512e-5, 8.031473407754675e-5, -3.937935980502516e-5, 2.2796950361225754e-5, -3.125887888018042e-4]
     >,
     "gamma" => #Nx.Tensor<
       f32[24]
       [0.6092047691345215, 0.47812265157699585, 0.693821370601654, 0.45786911249160767, 0.5387675166130066, 0.8261818289756775, 0.5807712078094482, 0.6510066986083984, 0.45650821924209595, 0.7799047827720642, 0.6193627715110779, 0.5857614874839783, 0.4286927878856659, 0.6337913870811462, 0.6911339163780212, 0.4652976393699646, 0.5756664872169495, 0.46622365713119507, 0.5369606018066406, 0.5942122340202332, 0.7727246880531311, 0.6141138076782227, 0.46364495158195496, 0.6175265908241272]
     >,
     "mean" => #Nx.Tensor<
       f32[24]
       [0.2229052484035492, -0.6143520474433899, -0.5115975737571716, 0.1679452657699585, -0.8705337047576904, -0.3773384094238281, 0.3830886483192444, 1.2518717050552368, -0.21099916100502014, 0.8486464023590088, 0.09219148010015488, -1.6454187631607056, 0.5216832160949707, -0.17828574776649475, 0.8979193568229675, -0.14078673720359802, 0.3999742567539215, 0.210261732339859, -1.1334097385406494, -0.07194671034812927, -0.5004290342330933, -1.782501459121704, -0.5524135231971741, -0.08081304281949997]
     >,
     "var" => #Nx.Tensor<
       f32[24]
       [0.14570137858390808, 0.0979502722620964, 0.20365412533283234, 0.08008237928152084, 0.09933093935251236, 0.3267607092857361, 0.0919615775346756, 0.2013940066099167, 0.08515820652246475, 0.30229872465133667, 0.1441039890050888, 0.13721634447574615, 0.06494998931884766, 0.37340933084487915, 0.24072226881980896, 0.06375624239444733, 0.13414347171783447, 0.09319230169057846, 0.13606417179107666, 0.10260697454214096, 0.2731955051422119, 0.2157246172428131, 0.07676148414611816, 0.1570962816476822]
     >
   },
   "mobilenetv20_features_linearbottleneck13_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[576]
       [-0.03612425923347473, 0.2621622085571289, 0.24970613420009613, 0.13786442577838898, 0.01677604205906391, -0.01019324455410242, 0.02858031913638115, 0.06309982389211655, 0.445175439119339, 0.12270592153072357, 0.37022992968559265, 0.10278589278459549, 0.05400807037949562, 0.2009139209985733, 0.2300354391336441, -0.056037537753582, 0.2122475951910019, 0.16997119784355164, 0.2927708923816681, -0.016652001067996025, -0.011880684643983841, 0.03291282057762146, 0.01846851408481598, 0.3031168282032013, 0.3798549771308899, -0.01126103289425373, 0.20447196066379547, 0.17122939229011536, -0.0733051523566246, 0.2149120718240738, 0.12058312445878983, -0.031003933399915695, 0.41413551568984985, 0.4456107020378113, 0.24590632319450378, -0.0205529797822237, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[576]
       [0.28325164318084717, 0.08031681180000305, 0.09241926670074463, 0.24785219132900238, 0.12529253959655762, 0.05622509494423866, 0.1654898077249527, 0.0985732153058052, 0.1314350813627243, -0.14117099344730377, 0.20672276616096497, 0.15235687792301178, 0.17320455610752106, 0.17095758020877838, 0.11314506083726883, 0.29315903782844543, 0.10191123932600021, 0.0798114538192749, 0.08920325338840485, 0.14326541125774384, 0.08023156970739365, 0.09082324057817459, 0.20910467207431793, 0.20261237025260925, 0.09189891070127487, 0.09407437592744827, 0.09949570894241333, 0.10999608039855957, 0.2922479808330536, 0.16595004498958588, 0.15531058609485626, 0.16434098780155182, 0.14855611324310303, 0.13777148723602295, 0.1618027538061142, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[576]
       [-0.003467562375590205, -0.03519834578037262, 0.006473586428910494, 0.0950964167714119, 0.005609882529824972, -1.5585363877831764e-16, 0.05394518002867699, 0.05065878853201866, -0.11292003840208054, 0.09376012533903122, -0.8194059729576111, -0.23955807089805603, 0.10773581266403198, -0.46938908100128174, 0.01830783486366272, -0.0031353626400232315, -0.044674184173345566, -0.09129905700683594, 0.05142587795853615, -0.0025106987450271845, -2.3497910297010094e-4, 0.02315840870141983, 0.009709623642265797, -0.3128106892108917, -0.059768788516521454, -1.0070249345517368e-6, -0.11782798171043396, -0.09451183676719666, -0.005958849564194679, -0.10387786477804184, -0.13271105289459229, -0.001979420892894268, -0.7967073917388916, -0.09457794576883316, ...]
     >,
     "var" => #Nx.Tensor<
       f32[576]
       [0.0010279034031555057, 0.00368467322550714, 0.0025712454225867987, 0.052593719214200974, 0.0019253796199336648, 1.0014506104153933e-18, 0.02545272558927536, 0.004488298203796148, 0.011829677037894726, 0.0027519832365214825, 0.2799582779407501, 0.06149482727050781, 0.004806542303413153, 0.0466470867395401, 0.0010941256769001484, 3.060759336221963e-4, 0.0022727937903255224, 0.008207143284380436, 0.0036502378061413765, 5.425590788945556e-4, 9.883774509944487e-6, 0.0035530165769159794, 0.01653285324573517, 0.007021782919764519, 0.013985850848257542, 1.463058829642705e-8, 0.014402402564883232, 0.004008142277598381, 5.924252327531576e-4, 0.007635928690433502, 0.008215081878006458, 1.4169601490721107e-4, 0.07166615128517151, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck7_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[384]
       [-0.05287745222449303, 0.03939635679125786, 0.008943471126258373, 0.016773579642176628, 0.024517597630620003, -0.2608976662158966, -0.013363493606448174, -0.04500827193260193, -0.0814104974269867, -0.14708103239536285, -0.1051139384508133, -0.03847385570406914, -0.029648277908563614, 0.07605069130659103, -0.042893584817647934, -0.03939994052052498, -0.29128172993659973, -0.03861434757709503, -0.02712533250451088, -0.019969439134001732, -0.07780560106039047, -0.02379690855741501, -0.1076209619641304, -0.028538363054394722, -0.1733534038066864, 0.02688787877559662, -0.05968484282493591, -0.01906135492026806, -0.07085005193948746, -0.04659796133637428, -0.10276751965284348, -0.10676652938127518, -0.05433874577283859, -0.08442538976669312, 0.3125787079334259, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[384]
       [0.19615307450294495, 0.13756610453128815, 0.1616729497909546, 0.16233323514461517, 0.1569656878709793, 0.30160748958587646, 0.1933916211128235, 0.18154680728912354, 0.24717095494270325, 0.29326149821281433, 0.2486879527568817, 0.12994907796382904, 0.24061599373817444, 0.13928817212581635, 0.1792965978384018, 0.23451970517635345, 0.23798415064811707, 0.21801051497459412, 0.15246279537677765, 0.20245827734470367, 0.21029138565063477, 0.20429369807243347, 0.26336348056793213, 0.2377043068408966, 0.1920943260192871, 0.14207608997821808, 0.2790174186229706, 0.22336551547050476, 0.2147824764251709, 0.18369099497795105, 0.27606746554374695, 0.28400060534477234, 0.16117443144321442, 0.18702644109725952, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[384]
       [-0.009776907041668892, 0.005012216046452522, -0.047731850296258926, 0.005270622204989195, 0.00557712372392416, -0.00950763002038002, -0.07496191561222076, 0.018082473427057266, 0.0058733015321195126, 0.019821494817733765, 0.008286324329674244, 0.006341803818941116, 0.007210697513073683, -0.021940631791949272, 0.005796361714601517, -0.008641923777759075, 0.012914237566292286, -5.078607355244458e-4, 0.021764321252703667, -0.0017188549973070621, 0.024174479767680168, -0.04216955602169037, -0.041497547179460526, -9.607977117411792e-4, 0.006745839957147837, 0.009343116544187069, 5.745989037677646e-4, 0.005417129956185818, 0.0023382019717246294, 0.002557831583544612, -0.03232024237513542, -0.06076114997267723, 0.0038470544386655092, ...]
     >,
     "var" => #Nx.Tensor<
       f32[384]
       [9.27670334931463e-4, 2.7918731211684644e-4, 4.3098966125398874e-4, 0.0011215894483029842, 0.0012489479267969728, 8.453890914097428e-4, 9.005925385281444e-4, 4.5995140681043267e-4, 9.497502469457686e-4, 7.083975942805409e-4, 6.839227862656116e-4, 9.455660620005801e-5, 8.202312747016549e-4, 4.4565973803400993e-4, 3.3621577313169837e-4, 0.0014562587020918727, 6.57685799524188e-4, 7.828701636753976e-4, 2.743843651842326e-4, 0.001106554176658392, 4.3269505840726197e-4, 5.598529824055731e-4, 6.088317022658885e-4, 0.0010905259987339377, 2.4314268375746906e-4, 3.704722039401531e-4, 0.0010972988093271852, 7.521603838540614e-4, 5.540369893424213e-4, 3.005551116075367e-4, 7.429083343595266e-4, 7.382072508335114e-4, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck8_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[384]
       [0.03583831712603569, 0.08322400599718094, 0.08347140997648239, 0.16407261788845062, -0.036046866327524185, -0.07233322411775589, -0.007247522007673979, 0.22039948403835297, 0.16788388788700104, 0.011028931476175785, -0.052548106759786606, 0.09673431515693665, -0.009311914443969727, 0.04754066839814186, -0.13346582651138306, 0.0462525449693203, 0.14666330814361572, 0.13988535106182098, 0.08048173785209656, 0.12783612310886383, 0.10406619310379028, 0.061243198812007904, 0.12299603223800659, 0.14473021030426025, 0.15518461167812347, -0.014591885730624199, 0.056073401123285294, 0.1325678676366806, -0.06322228908538818, -0.07645687460899353, -0.04285341501235962, -0.10308938473463058, 0.03268742933869362, -0.04713943228125572, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[384]
       [0.12965284287929535, 0.08416002988815308, 0.1387501209974289, 0.0941140204668045, 0.16713306307792664, 0.15179668366909027, 0.1911296546459198, 0.08323527127504349, 0.09111560136079788, 0.142413929104805, 0.12192691117525101, 0.10468165576457977, 0.16143439710140228, 0.14601951837539673, 0.16359323263168335, 0.13607709109783173, 0.0671626552939415, 0.09158993512392044, 0.11166514456272125, 0.12548623979091644, 0.08543179929256439, 0.11259179562330246, 0.06824684143066406, 0.04742749035358429, 0.07571908831596375, 0.129663348197937, 0.09579146653413773, 0.05607251077890396, 0.1601036787033081, 0.15332894027233124, 0.12319595366716385, 0.12570618093013763, 0.11699129641056061, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[384]
       [-2.3525810684077442e-4, 4.89452622787212e-6, -1.9739160779863596e-4, -3.639425849542022e-4, 1.6747551853768528e-4, 4.836112493649125e-4, -8.706151857040823e-5, -3.2933344482444227e-4, -1.4665452181361616e-4, 2.6016379706561565e-4, -7.966592238517478e-5, 3.506306529743597e-5, 1.9936829630751163e-4, 1.982947433134541e-4, -1.9341196457389742e-4, 6.499576265923679e-4, 9.932630928233266e-5, 2.0891131134703755e-4, 1.7995496455114335e-4, 1.8568814266473055e-4, 7.022052159300074e-5, -1.6939605120569468e-4, 1.3075539754936472e-5, 3.1177018536254764e-4, 2.119196142302826e-5, 2.907874295488e-4, 2.8831884264945984e-4, 2.1954343537800014e-4, 8.599817374488339e-5, 2.881138352677226e-4, 5.481694242917001e-4, 3.153556608594954e-4, ...]
     >,
     "var" => #Nx.Tensor<
       f32[384]
       [0.04813230410218239, 0.029002580791711807, 0.04511406272649765, 0.05107865855097771, 0.0974809005856514, 0.0741136446595192, 0.17005448043346405, 0.0427735410630703, 0.04020823910832405, 0.06400398910045624, 0.054310500621795654, 0.0434776172041893, 0.08294457197189331, 0.0641767680644989, 0.09278614073991776, 0.11118675768375397, 0.038928885012865067, 0.03660726547241211, 0.062162626534700394, 0.08595701307058334, 0.03366314619779587, 0.04004068672657013, 0.04966263100504875, 0.052451543509960175, 0.04023684188723564, 0.03982854261994362, 0.04733984172344208, 0.049878451973199844, 0.06715817749500275, 0.08148116618394852, 0.04592667892575264, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck13_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[576]
       [-0.13610589504241943, 0.029707221314311028, 0.07764720171689987, -0.0930335521697998, -0.06152098625898361, -0.0342678464949131, -0.05844024196267128, -0.002459034090861678, 0.08498657494783401, 0.1225726529955864, 0.17613333463668823, 0.1668761670589447, 0.1025225967168808, 0.21359726786613464, 0.07739955931901932, -0.141232430934906, 0.10692654550075531, 0.09533322602510452, 0.1416633278131485, -0.13436390459537506, -0.04763894900679588, -0.08429068326950073, 0.009572437964379787, 0.3033914864063263, 0.0484132245182991, -0.08781138062477112, 0.08987420797348022, 0.10349894315004349, -0.11735771596431732, 0.22785161435604095, 0.14756335318088531, -0.08500344306230545, 0.2637306749820709, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[576]
       [0.12385702133178711, 0.09890380501747131, 0.0771060660481453, 0.23227904736995697, 0.11977005004882812, 0.005590819288045168, 0.13258837163448334, 0.11466257274150848, 0.1158096119761467, 0.07918184250593185, 0.1444852203130722, 0.1032688096165657, 0.05017390102148056, 0.09281343966722488, 0.08032833784818649, 0.10454760491847992, 0.1351696252822876, 0.11343558132648468, 0.04775742068886757, 0.09771032631397247, 0.02657795138657093, 0.13788972795009613, 0.11856312304735184, 0.07200472801923752, 0.12148887664079666, 0.013946599327027798, 0.12394212186336517, 0.06285671144723892, 0.11612718552350998, 0.05600180849432945, 0.10418163239955902, 0.05338305979967117, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[576]
       [6.827032734690874e-7, 1.5534817521256628e-7, -5.093219215268618e-7, 4.983389771950897e-7, 1.2705514791377936e-6, 1.1588760173708579e-7, -1.417142243553826e-7, 1.1870022831317328e-7, 7.192203952399723e-7, -8.536806603842706e-7, 1.9767255707847653e-6, 1.2823227280023275e-6, 7.469633374057594e-7, 1.4591352055504103e-6, -1.3462655488183373e-6, 1.617389671082492e-7, -7.737047269529285e-8, -8.04814064281345e-8, 1.6728367313589843e-7, -1.2480612987531003e-8, 2.5469088882346114e-7, 3.542891704455542e-7, -2.3255206826888752e-7, -1.7710331121634226e-6, 7.398209618258988e-7, 2.879962437418726e-7, 1.0903487464020145e-6, 8.88873159965442e-7, 6.44339706923347e-7, 5.017034254706232e-7, -7.835358815100335e-7, ...]
     >,
     "var" => #Nx.Tensor<
       f32[576]
       [0.07495011389255524, 0.0859481543302536, 0.06548640131950378, 0.48155632615089417, 0.12509563565254211, 0.0017050545429810882, 0.1484297662973404, 0.11226049065589905, 0.14026911556720734, 0.084726482629776, 0.5485780239105225, 0.12343145161867142, 0.08799819648265839, 0.1470877230167389, 0.06761230528354645, 0.05614602938294411, 0.1670052707195282, 0.16908922791481018, 0.0753912553191185, 0.04617440328001976, 0.0031118346378207207, 0.08412214368581772, 0.1583808809518814, 0.19899487495422363, 0.11616887152194977, 0.008272404782474041, 0.11385060846805573, 0.059652168303728104, 0.06796986609697342, 0.1069282665848732, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck1_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[96]
       [0.3692082166671753, 0.019333837553858757, 0.74114990234375, 0.2654167115688324, 0.24005535244941711, 0.04846649616956711, 0.5354508757591248, 0.03959512710571289, 0.2306070178747177, 0.19200822710990906, -0.48113152384757996, 0.006648355163633823, -0.023521361872553825, 0.047236256301403046, -0.015566157177090645, 0.11580895632505417, -0.01523163728415966, -0.0030626191291958094, 0.3765759766101837, -0.01797349564731121, 0.15314938127994537, 0.18993273377418518, 0.3641800582408905, 0.4655008018016815, 0.22646495699882507, 0.0988355278968811, 0.24557138979434967, 0.3279559314250946, 0.0678282305598259, 0.271154522895813, 0.27961465716362, 0.01484939269721508, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[96]
       [0.1144828125834465, 0.29103899002075195, 0.19599959254264832, 0.11088918149471283, 0.11658263951539993, 0.2404186874628067, 0.05210668221116066, 0.279130220413208, 0.16053666174411774, 0.10330193489789963, 0.17037393152713776, 0.42853739857673645, 0.2756047546863556, 0.3480927646160126, 0.41298505663871765, 0.3465864956378937, 0.2918163239955902, 0.2078855335712433, 0.19570542871952057, 0.3232663571834564, 0.20460295677185059, 0.21620818972587585, 0.06636220216751099, 0.16237126290798187, 0.11230707168579102, 0.20659472048282623, 0.04134644940495491, 0.16707743704319, 0.18970169126987457, 0.08170202374458313, -0.14146141707897186, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[96]
       [7.715559913776815e-4, -3.0534371035173535e-4, -9.66210791375488e-5, -7.018522592261434e-5, 2.5279350666096434e-5, -3.164849185850471e-4, 1.1111465573776513e-4, 1.5042595623526722e-4, 4.5005575520917773e-4, -4.403230268508196e-4, -1.6362588212359697e-4, 1.3409779057838023e-4, 7.232923235278577e-5, 5.355612374842167e-4, 0.0017661581514403224, -4.008736868854612e-4, -1.3021935046708677e-5, 7.276511169038713e-4, 4.067765548825264e-4, -2.802436938509345e-5, 1.5298827202059329e-4, 1.1257633013883606e-4, -1.6121388762257993e-4, 4.559751832857728e-4, 1.0374663543188944e-4, -6.660629878751934e-5, 1.5707196143921465e-4, 3.371218917891383e-5, -4.4229160994291306e-4, -7.242524588946253e-5, ...]
     >,
     "var" => #Nx.Tensor<
       f32[96]
       [0.15688009560108185, 0.15477390587329865, 0.365114688873291, 0.15238657593727112, 0.16493624448776245, 0.09593218564987183, 0.17790137231349945, 0.19072653353214264, 0.10221002995967865, 0.06285464763641357, 0.13617345690727234, 0.4062573313713074, 0.16529865562915802, 0.3744926452636719, 0.36751073598861694, 0.3471532464027405, 0.08051957935094833, 0.1254158318042755, 0.2987021505832672, 0.11974798142910004, 0.17331230640411377, 0.1749797761440277, 0.12097641825675964, 0.2704963684082031, 0.1457628309726715, 0.14666584134101868, 0.099171482026577, 0.24007830023765564, 0.11633463203907013, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck11_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[576]
       [0.1379205733537674, 0.1852763593196869, 0.07007614523172379, 0.14786551892757416, 0.08499550819396973, 0.022980524227023125, 0.10445888340473175, -0.03380206599831581, 0.05621454492211342, 0.13879631459712982, 0.1278844028711319, 0.11503273248672485, 0.008648642338812351, 0.07986151427030563, 0.05827014893293381, 0.02635921724140644, 0.16676078736782074, 0.1696624606847763, -0.0049711973406374454, 0.11680029332637787, 0.09153503179550171, -0.09545790404081345, 0.10052905976772308, 0.2198638916015625, 0.09066508710384369, 0.021357741206884384, 0.1243291050195694, 0.13527752459049225, 0.09010029584169388, 0.07285095006227493, 0.09204164147377014, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[576]
       [0.1116461232304573, 0.08017633110284805, 0.11955763399600983, 0.09390335530042648, 0.09855279326438904, 0.14738331735134125, 0.0888434574007988, 0.13319958746433258, 0.10988157987594604, 0.09758485853672028, 0.10826665163040161, 0.1016610860824585, 0.09888888150453568, 0.12156049907207489, 0.11694491654634476, 0.119488425552845, 0.07615012675523758, 0.06269187480211258, 0.15567438304424286, 0.1064029410481453, 0.11337470263242722, 0.14443783462047577, 0.0840112492442131, 0.09276950359344482, 0.16790293157100677, 0.11844969540834427, 0.0930234044790268, 0.11310834437608719, 0.1162751242518425, 0.1104091927409172, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[576]
       [3.63793503765919e-7, 1.209602089602413e-7, 5.4800288751266635e-8, 2.552324076532386e-7, 3.2612440747925575e-8, 1.4689121030642127e-8, 1.1114344999896275e-7, -8.35849434110969e-8, 2.2601201976613083e-7, 1.6091770760340296e-7, 3.1186812066152925e-7, -1.1275391287313141e-8, -2.748528871165945e-8, -1.7028881416081276e-7, -1.3741856719207135e-7, -1.3034306789450056e-7, 4.320734205975896e-7, 5.223997732173302e-7, 1.3161741918565895e-8, 1.749855726984606e-7, 6.074458980265263e-8, 4.904541128780693e-7, 2.8845341670802327e-9, 3.752578692228781e-8, -8.265797646345163e-8, -2.6385482243540537e-8, 6.387549156272598e-9, -3.501408230022207e-8, 3.437627640323626e-7, ...]
     >,
     "var" => #Nx.Tensor<
       f32[576]
       [0.02586364559829235, 0.04419300705194473, 0.02695547789335251, 0.027057670056819916, 0.023040946573019028, 0.03236440196633339, 0.021965283900499344, 0.03861037641763687, 0.02072414942085743, 0.020027386024594307, 0.017105598002672195, 0.027752427384257317, 0.025669051334261894, 0.030634647235274315, 0.02309420332312584, 0.026124022901058197, 0.009496251121163368, 0.02531951665878296, 0.04465075209736824, 0.03168340027332306, 0.03208788484334946, 0.03684369474649429, 0.01707514189183712, 0.033794183284044266, 0.05672197788953781, 0.01841931976377964, 0.031393010169267654, 0.02732437662780285, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck15_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[960]
       [-0.08277615904808044, 0.12612955272197723, -0.007448040880262852, -0.04266936331987381, 0.07737076282501221, 0.0038264743052423, -0.04381075128912926, -0.061556603759527206, 0.07451092451810837, 0.00787589605897665, -0.020372798666357994, 0.03865779936313629, -0.05657007917761803, -0.048030927777290344, -0.07944952696561813, 0.07044878602027893, 0.009796769358217716, -0.0038760113529860973, -0.06943286210298538, 0.035002585500478745, -0.06959164142608643, -0.08562169969081879, -0.1489018350839615, -0.0025909075047820807, -0.01582941599190235, -0.0698612853884697, 0.05639493837952614, 0.3489030599594116, -0.032863155007362366, 0.07198280841112137, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[960]
       [0.3150080740451813, 0.16160078346729279, 0.22753824293613434, 0.23682628571987152, 0.16329334676265717, 0.1377934366464615, 0.19716574251651764, 0.15468646585941315, 0.14866210520267487, 0.1379021555185318, 0.18138086795806885, 0.1989935338497162, 0.2044813334941864, 0.22123073041439056, 0.22492459416389465, 0.158174529671669, 0.14002622663974762, 0.13364093005657196, 0.19789022207260132, 0.17698919773101807, 0.1932896226644516, 0.17972171306610107, 0.24500404298305511, 0.1422407180070877, 0.2906820476055145, 0.21549536287784576, 0.16308309137821198, 0.116135373711586, 0.22621788084506989, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[960]
       [0.006607928778976202, -0.005489333998411894, 0.001202044659294188, 0.016741983592510223, -0.021014627069234848, 3.9149439544416964e-4, -0.001380614354275167, 0.006143252365291119, -0.022377200424671173, 3.921858442481607e-4, 0.011808911338448524, -0.038100697100162506, -0.0027107521891593933, -0.002684985054656863, -0.003361391369253397, -0.027487538754940033, 0.0032275188714265823, 7.377289002761245e-4, -0.003101653652265668, -0.02393931709229946, -0.0026154580991715193, 4.0451946551911533e-4, 0.019924011081457138, 0.0014483053237199783, 0.012146278284490108, -6.006719195283949e-4, -0.02536875568330288, -0.005276186391711235, ...]
     >,
     "var" => #Nx.Tensor<
       f32[960]
       [6.60044839605689e-4, 2.093793300446123e-4, 6.339335814118385e-4, 8.293294231407344e-4, 4.367101937532425e-4, 3.0174001949490048e-5, 1.974338520085439e-4, 5.38308268005494e-5, 5.038247327320278e-4, 1.9694401999004185e-4, 3.8651813520118594e-4, 5.852532922290266e-4, 2.917999809142202e-4, 3.8506241980940104e-4, 3.85848106816411e-4, 4.808907106053084e-4, 8.100475679384544e-5, 6.1370710682240315e-6, 1.8867437029257417e-4, 3.83470905944705e-4, 5.22851652931422e-4, 4.0124301449395716e-5, 4.517717461567372e-4, 7.429129254887812e-6, 0.0013210318284109235, 2.3677274293731898e-4, 4.158295632805675e-4, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck9_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][384][1][1]
       [
         [
           [
             [-0.029621589928865433]
           ],
           [
             [0.03196175768971443]
           ],
           [
             [-0.04775437340140343]
           ],
           [
             [-0.02047593705356121]
           ],
           [
             [0.022681893780827522]
           ],
           [
             [0.0061656758189201355]
           ],
           [
             [-0.020124075934290886]
           ],
           [
             [-0.0419141948223114]
           ],
           [
             [-0.0958712100982666]
           ],
           [
             [-0.03112773969769478]
           ],
           [
             [0.033005211502313614]
           ],
           [
             [-0.0796334370970726]
           ],
           [
             [-0.08826649934053421]
           ],
           [
             [0.008132039569318295]
           ],
           [
             [0.13005898892879486]
           ],
           [
             [0.01112390961498022]
           ],
           [
             [-0.025056127458810806]
           ],
           [
             [-0.021656425669789314]
           ],
           [
             [-0.009619093500077724]
           ],
           [
             [-0.09488093852996826]
           ],
           [
             [-0.03190577030181885]
           ],
           [
             [0.08531086146831512]
           ],
           [
             [0.047572940587997437]
           ],
           [
             [-0.036207523196935654]
           ],
           [
             [-0.1315651535987854]
           ],
           [
             [-0.03746785223484039]
           ],
           [
             [-0.04043968394398689]
           ],
           [
             [-0.1008371040225029]
           ],
           [
             [-7.86909949965775e-4]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][576][1][1]
       [
         [
           [
             [-0.04216459393501282]
           ],
           [
             [-0.03359765186905861]
           ],
           [
             [0.021302534267306328]
           ],
           [
             [0.06966963410377502]
           ],
           [
             [-0.023086098954081535]
           ],
           [
             [0.02423691935837269]
           ],
           [
             [-0.0027191585395485163]
           ],
           [
             [-0.12812955677509308]
           ],
           [
             [0.10142464935779572]
           ],
           [
             [0.04367773234844208]
           ],
           [
             [-0.04808792844414711]
           ],
           [
             [-0.059124793857336044]
           ],
           [
             [-0.03544275462627411]
           ],
           [
             [-0.10292188823223114]
           ],
           [
             [-0.15540222823619843]
           ],
           [
             [0.08993023633956909]
           ],
           [
             [-0.0018712839810177684]
           ],
           [
             [0.04499407112598419]
           ],
           [
             [0.03717062994837761]
           ],
           [
             [0.1088862419128418]
           ],
           [
             [0.03789613023400307]
           ],
           [
             [0.011467562057077885]
           ],
           [
             [0.022153226658701897]
           ],
           [
             [0.011087901890277863]
           ],
           [
             [0.10830669850111008]
           ],
           [
             [-0.039926622062921524]
           ],
           [
             [0.03126271069049835]
           ],
           [
             [0.11330234259366989]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck2_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[144]
       [0.06266669183969498, 0.34174367785453796, 0.014498202130198479, 0.09171395748853683, -0.01248596515506506, 0.09248802065849304, 0.03913332521915436, -0.12479757517576218, 0.04716098681092262, 0.2623796761035919, 0.00875355675816536, 0.3866884410381317, 0.34793952107429504, 0.07759835571050644, 0.28590700030326843, 0.34494516253471375, 0.1602264642715454, 0.3233429491519928, -0.08500878512859344, 0.017791876569390297, 0.012547487393021584, 0.07443935424089432, 0.470037043094635, 0.06998748332262039, -0.1145705133676529, 0.35432684421539307, 0.0036894683726131916, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[144]
       [0.14569389820098877, 0.08483350276947021, 0.18709848821163177, 0.08500487357378006, 0.0695008710026741, 0.13266052305698395, 0.23066304624080658, 0.1294335573911667, 0.15951289236545563, 0.1612233966588974, 0.04986649379134178, 0.24220621585845947, 0.13059896230697632, 0.2140040248632431, 0.16401590406894684, 0.04375945031642914, 0.1284344643354416, 0.10862768441438675, 0.22237107157707214, 0.20400500297546387, 0.2341395914554596, 0.36944401264190674, 0.1626269817352295, 0.1425088793039322, 0.37626516819000244, 0.10176140815019608, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[144]
       [0.04842672124505043, 0.011466434225440025, -0.28272300958633423, 0.04086904227733612, -5.605193857299268e-45, 0.03728378936648369, -0.07810080051422119, -0.0095414062961936, 0.019914837554097176, -0.26291584968566895, 1.200515914815965e-22, 0.4378272294998169, -0.07176417112350464, -0.20427647233009338, -0.40337827801704407, 0.07497201859951019, 0.06511131674051285, 0.02972135879099369, -0.029825769364833832, 0.03816595301032066, 0.008758057840168476, -0.5172790288925171, 0.18057824671268463, 0.03219117224216461, 0.25589102506637573, ...]
     >,
     "var" => #Nx.Tensor<
       f32[144]
       [0.0016082772053778172, 0.008874273858964443, 0.014802338555455208, 0.010124930180609226, 5.605193857299268e-45, 0.0027794779743999243, 0.019753890112042427, 6.709994049742818e-4, 0.004814377520233393, 0.018505984917283058, 3.875728477835647e-26, 0.024166392162442207, 0.011687243357300758, 0.018400687724351883, 0.00805666297674179, 0.003925488330423832, 0.004267431795597076, 0.005561113357543945, 0.00586837250739336, 0.01074590440839529, 0.023998573422431946, 0.010233975946903229, 0.007847330532968044, 0.0024662334471940994, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck3_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[144][1][3][3]
       [
         [
           [
             [0.007945280522108078, -0.5135775208473206, 0.642527163028717],
             [0.09571424126625061, -0.4191504120826721, -0.42046400904655457],
             [0.2738749086856842, -1.341597080230713, 1.7941224575042725]
           ]
         ],
         [
           [
             [0.36698344349861145, 0.8006415963172913, 1.0355921983718872],
             [0.2570298910140991, 0.9058852791786194, 0.1435706615447998],
             [0.08509054780006409, -0.29565608501434326, -0.19910597801208496]
           ]
         ],
         [
           [
             [-0.07826036959886551, 1.5095280408859253, -0.7258926033973694],
             [-0.1067059189081192, -0.07988816499710083, -0.4704415500164032],
             [-2.0861942768096924, 0.03404847905039787, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[384]
       [-0.14200253784656525, -0.08138399571180344, 0.020918697118759155, -0.030422303825616837, -0.054856158792972565, 0.0021326385904103518, -0.18490885198116302, -0.03978404402732849, 0.17561353743076324, 0.007802810054272413, -0.05954292416572571, 0.001685986528173089, -0.021205948665738106, 0.09003060311079025, 0.03726910054683685, 0.14319199323654175, 0.15906931459903717, 0.09075694531202316, 0.14122611284255981, 0.19317775964736938, 0.11409415304660797, 0.2862238883972168, 0.21863330900669098, 0.31881415843963623, 0.15870404243469238, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[384]
       [0.08955207467079163, 0.12251311540603638, 0.13474930822849274, 0.12169622629880905, 0.10408146679401398, 0.12673817574977875, 0.09429759532213211, 0.004978157579898834, 0.06450905650854111, 0.12292398512363434, 0.008911419659852982, 0.16245806217193604, 0.12929747998714447, 0.07689117640256882, 0.13534259796142578, 0.0939403846859932, 0.07975732535123825, 0.07448924332857132, 0.02857014909386635, 0.05634821206331253, 0.04961058497428894, 0.08511262387037277, 0.1229783222079277, 0.09482692182064056, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[384]
       [2.3076812794897705e-4, 8.846540004014969e-4, 2.087028551613912e-4, 2.626747591421008e-4, -1.6135886835400015e-4, 3.3246324164792895e-4, 1.477891200920567e-4, 8.788402919890359e-5, -5.301489727571607e-4, 3.370658669155091e-4, 1.6383620095439255e-4, 8.14800092484802e-4, 4.6927749644964933e-4, 2.0585938182193786e-4, 7.599047385156155e-4, -1.2353641795925796e-4, -2.2819229343440384e-4, 1.4898847439326346e-4, -1.760231316438876e-5, 1.1794252350227907e-4, 3.7943560164421797e-4, 8.82304084370844e-6, 8.040065877139568e-4, ...]
     >,
     "var" => #Nx.Tensor<
       f32[384]
       [0.03708014264702797, 0.1853957325220108, 0.18863092362880707, 0.12991510331630707, 0.08043266087770462, 0.15097717940807343, 0.07633852958679199, 0.003970237914472818, 0.38535448908805847, 0.13255122303962708, 0.01239611953496933, 0.37033578753471375, 0.1896938532590866, 0.13335467875003815, 0.267678827047348, 0.12042298913002014, 0.16346299648284912, 0.09761403501033783, 0.12983210384845734, 0.211727574467659, 0.12430556863546371, 0.3960045278072357, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck14_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[160]
       [-8.388164474126825e-8, -5.7530677111117257e-8, 7.027283999150313e-8, -9.3144599588868e-8, 4.638787132194011e-8, -9.98987204070545e-8, 1.4243799739688257e-7, -1.1660619492204205e-7, 6.76579361424956e-8, 1.2566950147174794e-7, -1.1262851984383815e-7, 6.708535948973804e-8, 2.0174957171548158e-7, 1.1858566040245933e-7, 3.9078233982081656e-8, -1.565037024420235e-7, 1.173001891174863e-7, -1.3883133931358316e-7, -1.3765874484761298e-7, -1.1415199452358138e-7, -3.374253765286994e-8, 1.203010810968408e-7, -1.509711751168652e-7, 6.394230922524002e-8, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[160]
       [0.18752147257328033, 0.34522393345832825, 0.24475780129432678, 0.33060014247894287, 0.21970151364803314, 0.3153970539569855, 0.28512927889823914, 0.169156014919281, 0.3166297674179077, 0.2279355376958847, 0.21180318295955658, 0.22759808599948883, 0.23080326616764069, 0.17287319898605347, 0.27370786666870117, 0.2255641222000122, 0.2651505172252655, 0.17854700982570648, 0.257531076669693, 0.2266509234905243, 0.2487022578716278, 0.22212600708007812, 0.19973145425319672, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[160]
       [-0.10289469361305237, 0.13421736657619476, 0.3127436637878418, -0.010098671540617943, 0.11501706391572952, 0.07170626521110535, -0.12338873744010925, -0.09662780910730362, 0.10249986499547958, 0.039654918015003204, 0.23144079744815826, -0.02762533910572529, -0.2323276251554489, -0.14802084863185883, -0.17405559122562408, -0.05928152799606323, 0.30801624059677124, 0.03448152542114258, 0.10294467210769653, -0.14233750104904175, 0.042918454855680466, -0.15366889536380768, ...]
     >,
     "var" => #Nx.Tensor<
       f32[160]
       [0.01982937380671501, 0.07223469018936157, 0.03559407591819763, 0.05254514887928963, 0.03399147093296051, 0.05125325173139572, 0.045940618962049484, 0.015112430788576603, 0.0665295273065567, 0.029054192826151848, 0.024110015481710434, 0.028089039027690887, 0.032291218638420105, 0.017801793292164803, 0.04934801161289215, 0.028767650946974754, 0.036390747874975204, 0.020782455801963806, 0.038260187953710556, 0.03192899003624916, 0.0354829803109169, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck0_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[16]
       [-0.0010560675291344523, -6.787155289202929e-4, -7.046025712043047e-4, -9.398604743182659e-4, -8.966192253865302e-4, -2.11737904464826e-4, -5.149749922566116e-4, -0.0010324727045372128, -7.925425306893885e-4, -9.090156527236104e-5, -4.453391011338681e-4, -9.805577574297786e-4, -7.45584664400667e-4, -7.640189724043012e-4, -0.0011470100143924356, -3.8889769348315895e-4]
     >,
     "gamma" => #Nx.Tensor<
       f32[16]
       [0.5097491145133972, 0.6192775368690491, 0.657399594783783, 0.5558041930198669, 0.6164559721946716, 0.6381960511207581, 0.6152222156524658, 0.4946329891681671, 0.5616724491119385, 0.6567564606666565, 0.5526584982872009, 0.42007970809936523, 0.4637764096260071, 0.5626768469810486, 0.6189385056495667, 0.6354565620422363]
     >,
     "mean" => #Nx.Tensor<
       f32[16]
       [0.020797766745090485, -0.4769341051578522, -0.4614068269729614, 1.7337007522583008, 0.3933316469192505, -0.8451346755027771, 1.245340347290039, -0.48207274079322815, 0.8380969166755676, -2.1938257217407227, 2.286595106124878, 0.04284362867474556, -1.1523691415786743, -0.37927713990211487, -3.048325300216675, -0.7846338748931885]
     >,
     "var" => #Nx.Tensor<
       f32[16]
       [0.020222526043653488, 0.055290378630161285, 0.06333210319280624, 0.026319650933146477, 0.03692027926445007, 0.0943082720041275, 0.044809967279434204, 0.05162627622485161, 0.03870456665754318, 0.13497591018676758, 0.04896015673875809, 0.0536123663187027, 0.03798253461718559, 0.045080412179231644, 0.05210200324654579, 0.049917370080947876]
     >
   },
   "mobilenetv20_features_linearbottleneck7_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][1][3][3]
       [
         [
           [
             [0.09069628268480301, 0.011556080542504787, -0.044405367225408554],
             [0.1974998563528061, -0.05144780874252319, -0.15406860411167145],
             [0.03294772282242775, -0.08055609464645386, -0.10401368886232376]
           ]
         ],
         [
           [
             [0.004530614707618952, -0.016135336831212044, 0.010837683454155922],
             [-0.013911036774516106, -0.17941752076148987, 0.0742872878909111],
             [0.03701423481106758, 0.17489252984523773, 0.049539677798748016]
           ]
         ],
         [
           [
             [-0.03442509472370148, -0.04338786378502846, -8.070958429016173e-4],
             [-0.06931383162736893, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck1_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][16][1][1]
       [
         [
           [
             [-0.26627230644226074]
           ],
           [
             [0.1223592683672905]
           ],
           [
             [-0.17675793170928955]
           ],
           [
             [-0.14175795018672943]
           ],
           [
             [-0.04847275838255882]
           ],
           [
             [0.02572016417980194]
           ],
           [
             [0.18247880041599274]
           ],
           [
             [-0.22272470593452454]
           ],
           [
             [-0.0375950001180172]
           ],
           [
             [-0.12036115676164627]
           ],
           [
             [-0.42806223034858704]
           ],
           [
             [0.2049194574356079]
           ],
           [
             [0.025416148826479912]
           ],
           [
             [-0.20868347585201263]
           ],
           [
             [-0.00869692862033844]
           ],
           [
             [0.10014521330595016]
           ]
         ],
         [
           [
             [0.1901395469903946]
           ],
           [
             [-0.15854698419570923]
           ],
           [
             [-0.05858785659074783]
           ],
           [
             [-0.016011832281947136]
           ],
           [
             [-0.07135247439146042]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck2_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[24]
       [-6.408389890566468e-5, -1.0128262510988861e-4, -4.3479518353706226e-5, -1.589354214956984e-4, -7.557999197160825e-5, -2.3196584152174182e-5, -2.216532493548584e-6, 9.525781933916733e-6, -8.600373985245824e-5, -2.3008385323919356e-4, -2.3174902889877558e-4, -2.012878976529464e-4, -9.838923870120198e-5, -1.22729703434743e-4, -9.646145917940885e-5, -2.797471643134486e-5, -1.039204653352499e-4, -1.2249314750079066e-4, -8.16335596027784e-5, -1.541685633128509e-4, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[24]
       [0.5065965056419373, 0.44786974787712097, 0.5534567832946777, 0.3781253695487976, 0.44674113392829895, 0.28593751788139343, 0.5029821395874023, 0.19650447368621826, 0.3918337821960449, 0.33050450682640076, 0.31709131598472595, 0.1402808576822281, 0.4198921024799347, 0.02611447498202324, 0.5337468981742859, 0.49816271662712097, 0.32386210560798645, 0.40961867570877075, 0.3363959491252899, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[24]
       [-0.49676281213760376, -0.5412184000015259, 0.2777172029018402, 0.30402812361717224, -0.038367077708244324, 0.19269412755966187, -0.20861922204494476, 0.3330950438976288, -0.35258007049560547, -0.0413212850689888, -0.44278526306152344, 0.43556883931159973, -0.007557356264442205, -0.049789782613515854, -0.1810888946056366, -0.1666330248117447, 0.15477608144283295, 0.1351827085018158, ...]
     >,
     "var" => #Nx.Tensor<
       f32[24]
       [0.0421639047563076, 0.03801903501152992, 0.04582618921995163, 0.028266852721571922, 0.025787167251110077, 0.026369646191596985, 0.03499341756105423, 0.02089923992753029, 0.02381058968603611, 0.04512317106127739, 0.018083490431308746, 0.02251904457807541, 0.023619014769792557, 0.005619029048830271, 0.04942920058965683, 0.02654757723212242, 0.02269740216434002, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck8_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][1][3][3]
       [
         [
           [
             [-0.0020327651873230934, -0.03829871118068695, -0.013563335873186588],
             [-0.07023981958627701, -0.0530773401260376, -0.05999664217233658],
             [-0.02172546088695526, -0.05160605534911156, -0.02340804785490036]
           ]
         ],
         [
           [
             [-0.06414368003606796, -0.09891236573457718, 0.09147416055202484],
             [-0.035863131284713745, -0.026917265728116035, 0.07736452668905258],
             [0.08244364708662033, 0.10789858549833298, -0.16784991323947906]
           ]
         ],
         [
           [
             [0.13550637662410736, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[576][1][3][3]
       [
         [
           [
             [0.0561985969543457, -0.08534234762191772, -0.0728735700249672],
             [0.20427849888801575, -0.03384411707520485, -0.18976956605911255],
             [0.10962294042110443, 0.006988390814512968, -0.03254517540335655]
           ]
         ],
         [
           [
             [-0.007281317375600338, -0.10424316674470901, -0.0011419359361752868],
             [-0.0659799799323082, -0.2147035449743271, -0.10787361860275269],
             [0.004407626111060381, 0.44921964406967163, 0.06901120394468307]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck14_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][160][1][1]
       [
         [
           [
             [0.037727002054452896]
           ],
           [
             [-0.07920914888381958]
           ],
           [
             [-0.0031863711774349213]
           ],
           [
             [-0.017933860421180725]
           ],
           [
             [-0.012083934620022774]
           ],
           [
             [-0.03305855393409729]
           ],
           [
             [0.010104630142450333]
           ],
           [
             [-0.01614421419799328]
           ],
           [
             [0.0415346585214138]
           ],
           [
             [-0.05910014733672142]
           ],
           [
             [3.380556881893426e-4]
           ],
           [
             [0.02735193446278572]
           ],
           [
             [-0.0788668617606163]
           ],
           [
             [-0.028101010248064995]
           ],
           [
             [-0.01531447283923626]
           ],
           [
             [-0.04052167013287544]
           ],
           [
             [0.0598626434803009]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck15_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][1][3][3]
       [
         [
           [
             [0.008491688407957554, 0.01882096938788891, -0.001548346015624702],
             [0.009763970039784908, 0.2359771430492401, 0.03868749365210533],
             [-0.05791253224015236, 0.028106149286031723, -0.07315897196531296]
           ]
         ],
         [
           [
             [-0.01212818082422018, -0.16292721033096313, -0.033156707882881165],
             [0.005929793696850538, 0.2561517059803009, 0.014003193937242031],
             [-0.04701530933380127, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck7_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][384][1][1]
       [
         [
           [
             [-0.03408721089363098]
           ],
           [
             [-0.017362290993332863]
           ],
           [
             [0.01995295286178589]
           ],
           [
             [-0.051705457270145416]
           ],
           [
             [0.0689198449254036]
           ],
           [
             [0.08685684949159622]
           ],
           [
             [-0.09933704882860184]
           ],
           [
             [0.07178854197263718]
           ],
           [
             [0.044526729732751846]
           ],
           [
             [-0.018590005114674568]
           ],
           [
             [0.013196618296205997]
           ],
           [
             [-0.0795169249176979]
           ],
           [
             [-0.10800802707672119]
           ],
           [
             [-0.042785294353961945]
           ],
           [
             [0.028526680544018745]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck4_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][1][3][3]
       [
         [
           [
             [-0.031621839851140976, 0.18290314078330994, 0.040609210729599],
             [-0.19228996336460114, 0.031516727060079575, 0.09460115432739258],
             [-0.026457147672772408, 0.05563095211982727, 0.0026122310664504766]
           ]
         ],
         [
           [
             [-0.00655049504712224, -0.004564541857689619, -0.01380461547523737],
             [-0.002132135210558772, -0.17951297760009766, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck16_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][160][1][1]
       [
         [
           [
             [0.04804954305291176]
           ],
           [
             [-0.0361880324780941]
           ],
           [
             [0.09869731217622757]
           ],
           [
             [0.03218498453497887]
           ],
           [
             [-0.05613848939538002]
           ],
           [
             [0.03235843405127525]
           ],
           [
             [0.035072293132543564]
           ],
           [
             [0.04535597562789917]
           ],
           [
             [0.031406670808792114]
           ],
           [
             [0.003971233498305082]
           ],
           [
             [0.016072509810328484]
           ],
           [
             [-0.041956160217523575]
           ],
           [
             [-0.018940964713692665]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[96]
       [2.1335902999908285e-7, 3.535939754328865e-7, 5.073925990473072e-7, 5.08670098042785e-7, 3.3390904263796983e-7, 4.946827516505437e-7, 4.795706445293035e-7, 3.743390948329761e-7, 5.272184466775798e-7, 5.93354513966915e-7, 2.1419212714590685e-7, 6.500247309304541e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[96]
       [0.2979993522167206, 0.3140060007572174, 0.2445983588695526, 0.2865447998046875, 0.35691216588020325, 0.3592449426651001, 0.27936986088752747, 0.25922414660453796, 0.2633695900440216, 0.2901868522167206, 0.2989289462566376, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[96]
       [-0.02239207923412323, -0.5143733024597168, 0.05194510519504547, -0.3102657198905945, -0.03491494804620743, 0.21602749824523926, 0.09913935512304306, -0.2323945015668869, -0.47182056307792664, 0.019973447546362877, ...]
     >,
     "var" => #Nx.Tensor<
       f32[96]
       [0.10519497841596603, 0.09767355769872665, 0.06126777082681656, 0.09396423399448395, 0.17368470132350922, 0.14595730602741241, 0.0813060998916626, 0.07312573492527008, 0.059541765600442886, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck14_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][1][3][3]
       [
         [
           [
             [0.03072543255984783, 0.024941517040133476, -0.044987138360738754],
             [0.2105334848165512, -0.09797866642475128, -0.12514935433864594],
             [0.049303051084280014, -0.004124515224248171, -0.0421803742647171]
           ]
         ],
         [
           [
             [0.027953263372182846, -0.0323273129761219, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck12_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[576][96][1][1]
       [
         [
           [
             [-0.020311787724494934]
           ],
           [
             [0.00342985475435853]
           ],
           [
             [-0.04346230626106262]
           ],
           [
             [-0.016626032069325447]
           ],
           [
             [-0.012114123441278934]
           ],
           [
             [-0.012785198166966438]
           ],
           [
             [0.056096695363521576]
           ],
           [
             [0.09021388739347458]
           ],
           [
             [0.16139942407608032]
           ],
           [
             [0.06402886658906937]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck0_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[32]
       [1.037217140197754, 1.2865849733352661, 0.7645959258079529, 0.6374938488006592, 1.069508671760559, 0.6837575435638428, 0.9550023674964905, 0.9839419722557068, 1.1442939043045044, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[32]
       [0.2217620462179184, 0.11614178121089935, 0.1586529165506363, 0.15199987590312958, 0.23966756463050842, 0.17114117741584778, 0.18509696424007416, 0.18513236939907074, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[32]
       [0.6259090900421143, -0.08340238034725189, 1.0731103420257568, 2.4088242053985596, -5.609257698059082, 3.4169440269470215, 7.146847724914551, ...]
     >,
     "var" => #Nx.Tensor<
       f32[32]
       [0.22989171743392944, 0.007035647518932819, 0.011941581964492798, 0.3442769944667816, 2.0481817722320557, 0.31869152188301086, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck6_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][1][3][3]
       [
         [
           [
             [0.014363789930939674, 0.13237811625003815, 0.09334706515073776],
             [0.12919868528842926, -0.35225632786750793, -0.1466313749551773],
             [0.06540388613939285, -0.08005677908658981, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck16_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[320]
       [-5.189255602999765e-7, -2.7314237627251714e-7, -2.1782213366350334e-7, 1.9330661871208576e-7, 2.6874428016299134e-8, -2.4903325623881756e-8, -2.696318688322208e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[320]
       [0.25196632742881775, 0.2486555427312851, 0.24830971658229828, 0.24681176245212555, 0.24941202998161316, 0.2485744059085846, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[320]
       [-0.13978023827075958, -0.05438295006752014, -0.14319643378257751, 0.1896054595708847, -0.13925714790821075, ...]
     >,
     "var" => #Nx.Tensor<
       f32[320]
       [0.008434775285422802, 0.007849316112697124, 0.0076168072409927845, 0.007443031296133995, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck10_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][384][1][1]
       [
         [
           [
             [0.01824713684618473]
           ],
           [
             [0.08710300177335739]
           ],
           [
             [-0.10191874951124191]
           ],
           [
             [-0.07710960507392883]
           ],
           [
             [-0.042343202978372574]
           ],
           [
             [-0.014866536483168602]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck8_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][64][1][1]
       [
         [
           [
             [0.015875784680247307]
           ],
           [
             [0.00680553587153554]
           ],
           [
             [0.09295158088207245]
           ],
           [
             [-0.05411294475197792]
           ],
           [
             [0.026204092428088188]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[32][3][3][3]
       [
         [
           [
             [-0.11137320101261139, -0.22142910957336426, 0.10717468708753586],
             [4.23799006966874e-4, ...],
             ...
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[96]
       [3.9526614159512974e-7, 5.761955890193349e-7, 7.01267651948001e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[96]
       [0.21182486414909363, 0.20167435705661774, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[96]
       [-0.039972804486751556, ...]
     >,
     "var" => #Nx.Tensor<
       f32[96]
       [...]
     >
   },
   "mobilenetv20_features_linearbottleneck6_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][192][1][1]
       [
         [
           [
             [0.07365027070045471]
           ],
           [
             [-0.03652043268084526]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck3_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[32][144][1][1]
       [
         [
           [
             [0.02765008620917797]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck4_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][32][1][1]
       [
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_conv0_fwd" => %{...},
   ...
 }}
input_template = Nx.template({1, 3, target_height, target_width}, :f32)
Axon.Display.as_graph(cnn_base, input_template)
graph TD;
33["mobilenetv20_features_linearbottleneck2_batchnorm1_fwd (:batch_norm) {1, 144, 40, 40}"];
168["mobilenetv20_features_linearbottleneck16_batchnorm2_fwd (:batch_norm) {1, 320, 5, 5}"];
117["mobilenetv20_features_linearbottleneck11_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
12["mobilenetv20_features_relu0_fwd (:relu) {1, 32, 80, 80}"];
157["mobilenetv20_features_linearbottleneck15_conv2_fwd (:conv) {1, 160, 5, 5}"];
132["add_7 (:add) {1, 96, 10, 10}"];
73["mobilenetv20_features_linearbottleneck6_conv2_fwd (:conv) {1, 64, 20, 20}"];
44["mobilenetv20_features_linearbottleneck3_relu1_fwd (:relu) {1, 144, 20, 20}"];
124["mobilenetv20_features_linearbottleneck12_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
170["mobilenetv20_features_batchnorm1_fwd (:batch_norm) {1, 1280, 5, 5}"];
23["mobilenetv20_features_linearbottleneck1_relu0_fwd (:relu) {1, 96, 80, 80}"];
29["mobilenetv20_features_linearbottleneck2_conv0_fwd (:conv) {1, 144, 40, 40}"];
47["mobilenetv20_features_linearbottleneck4_conv0_fwd (:conv) {1, 192, 20, 20}"];
89["mobilenetv20_features_linearbottleneck8_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
61["mobilenetv20_features_linearbottleneck5_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
30["mobilenetv20_features_linearbottleneck2_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
43["mobilenetv20_features_linearbottleneck3_batchnorm1_fwd (:batch_norm) {1, 144, 20, 20}"];
163["mobilenetv20_features_linearbottleneck16_relu0_fwd (:relu) {1, 960, 5, 5}"];
39["mobilenetv20_features_linearbottleneck3_conv0_fwd (:conv) {1, 144, 40, 40}"];
131["container_7 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
45["mobilenetv20_features_linearbottleneck3_conv2_fwd (:conv) {1, 32, 20, 20}"];
48["mobilenetv20_features_linearbottleneck4_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
145["mobilenetv20_features_linearbottleneck14_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
171["mobilenetv20_features_relu1_fwd (:relu) {1, 1280, 5, 5}"];
57["mobilenetv20_features_linearbottleneck5_conv0_fwd (:conv) {1, 192, 20, 20}"];
143["mobilenetv20_features_linearbottleneck14_relu0_fwd (:relu) {1, 960, 5, 5}"];
113["mobilenetv20_features_linearbottleneck11_conv0_fwd (:conv) {1, 576, 10, 10}"];
26["mobilenetv20_features_linearbottleneck1_relu1_fwd (:relu) {1, 96, 40, 40}"];
69["mobilenetv20_features_linearbottleneck6_relu0_fwd (:relu) {1, 192, 20, 20}"];
88["mobilenetv20_features_linearbottleneck8_conv1_fwd (:conv) {1, 384, 20, 20}"];
166["mobilenetv20_features_linearbottleneck16_relu1_fwd (:relu) {1, 960, 5, 5}"];
144["mobilenetv20_features_linearbottleneck14_conv1_fwd (:conv) {1, 960, 5, 5}"];
141["mobilenetv20_features_linearbottleneck14_conv0_fwd (:conv) {1, 960, 5, 5}"];
152["mobilenetv20_features_linearbottleneck15_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
63["mobilenetv20_features_linearbottleneck5_conv2_fwd (:conv) {1, 32, 20, 20}"];
165["mobilenetv20_features_linearbottleneck16_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
100["mobilenetv20_features_linearbottleneck9_relu1_fwd (:relu) {1, 384, 20, 20}"];
71["mobilenetv20_features_linearbottleneck6_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
133["mobilenetv20_features_linearbottleneck13_conv0_fwd (:conv) {1, 576, 10, 10}"];
46["mobilenetv20_features_linearbottleneck3_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
31["mobilenetv20_features_linearbottleneck2_relu0_fwd (:relu) {1, 144, 40, 40}"];
158["mobilenetv20_features_linearbottleneck15_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
167["mobilenetv20_features_linearbottleneck16_conv2_fwd (:conv) {1, 320, 5, 5}"];
98["mobilenetv20_features_linearbottleneck9_conv1_fwd (:conv) {1, 384, 20, 20}"];
81["mobilenetv20_features_linearbottleneck7_conv2_fwd (:conv) {1, 64, 20, 20}"];
11["mobilenetv20_features_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
37["container_0 (:container) {{1, 24, 40, 40}, {1, 24, 40, 40}}"];
153["mobilenetv20_features_linearbottleneck15_relu0_fwd (:relu) {1, 960, 5, 5}"];
9[/"data (:input) {1, 3, 160, 160}"/];
76["mobilenetv20_features_linearbottleneck7_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
173["mobilenetv20_output_pred_fwd (:conv) {1, 1000, 1, 1}"];
32["mobilenetv20_features_linearbottleneck2_conv1_fwd (:conv) {1, 144, 40, 40}"];
34["mobilenetv20_features_linearbottleneck2_relu1_fwd (:relu) {1, 144, 40, 40}"];
25["mobilenetv20_features_linearbottleneck1_batchnorm1_fwd (:batch_norm) {1, 96, 40, 40}"];
28["mobilenetv20_features_linearbottleneck1_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
85["mobilenetv20_features_linearbottleneck8_conv0_fwd (:conv) {1, 384, 20, 20}"];
91["mobilenetv20_features_linearbottleneck8_conv2_fwd (:conv) {1, 64, 20, 20}"];
120["mobilenetv20_features_linearbottleneck11_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
64["mobilenetv20_features_linearbottleneck5_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
146["mobilenetv20_features_linearbottleneck14_relu1_fwd (:relu) {1, 960, 5, 5}"];
109["mobilenetv20_features_linearbottleneck10_batchnorm1_fwd (:batch_norm) {1, 384, 10, 10}"];
72["mobilenetv20_features_linearbottleneck6_relu1_fwd (:relu) {1, 192, 20, 20}"];
110["mobilenetv20_features_linearbottleneck10_relu1_fwd (:relu) {1, 384, 10, 10}"];
68["mobilenetv20_features_linearbottleneck6_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
86["mobilenetv20_features_linearbottleneck8_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
116["mobilenetv20_features_linearbottleneck11_conv1_fwd (:conv) {1, 576, 10, 10}"];
51["mobilenetv20_features_linearbottleneck4_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
80["mobilenetv20_features_linearbottleneck7_relu1_fwd (:relu) {1, 384, 20, 20}"];
38["add_0 (:add) {1, 24, 40, 40}"];
137["mobilenetv20_features_linearbottleneck13_batchnorm1_fwd (:batch_norm) {1, 576, 5, 5}"];
126["mobilenetv20_features_linearbottleneck12_conv1_fwd (:conv) {1, 576, 10, 10}"];
13["mobilenetv20_features_linearbottleneck0_conv0_fwd (:conv) {1, 32, 80, 80}"];
59["mobilenetv20_features_linearbottleneck5_relu0_fwd (:relu) {1, 192, 20, 20}"];
40["mobilenetv20_features_linearbottleneck3_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
123["mobilenetv20_features_linearbottleneck12_conv0_fwd (:conv) {1, 576, 10, 10}"];
77["mobilenetv20_features_linearbottleneck7_relu0_fwd (:relu) {1, 384, 20, 20}"];
95["mobilenetv20_features_linearbottleneck9_conv0_fwd (:conv) {1, 384, 20, 20}"];
130["mobilenetv20_features_linearbottleneck12_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
41["mobilenetv20_features_linearbottleneck3_relu0_fwd (:relu) {1, 144, 40, 40}"];
20["mobilenetv20_features_linearbottleneck0_batchnorm2_fwd (:batch_norm) {1, 16, 80, 80}"];
78["mobilenetv20_features_linearbottleneck7_conv1_fwd (:conv) {1, 384, 20, 20}"];
174["mobilenetv20_output_flatten0_reshape0 (:reshape) {1000}"];
15["mobilenetv20_features_linearbottleneck0_relu0_fwd (:relu) {1, 32, 80, 80}"];
147["mobilenetv20_features_linearbottleneck14_conv2_fwd (:conv) {1, 160, 5, 5}"];
14["mobilenetv20_features_linearbottleneck0_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
108["mobilenetv20_features_linearbottleneck10_conv1_fwd (:conv) {1, 384, 10, 10}"];
60["mobilenetv20_features_linearbottleneck5_conv1_fwd (:conv) {1, 192, 20, 20}"];
162["mobilenetv20_features_linearbottleneck16_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
103["container_5 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
122["add_6 (:add) {1, 96, 10, 10}"];
140["mobilenetv20_features_linearbottleneck13_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
83["container_3 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
74["mobilenetv20_features_linearbottleneck6_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
139["mobilenetv20_features_linearbottleneck13_conv2_fwd (:conv) {1, 160, 5, 5}"];
99["mobilenetv20_features_linearbottleneck9_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
112["mobilenetv20_features_linearbottleneck10_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
97["mobilenetv20_features_linearbottleneck9_relu0_fwd (:relu) {1, 384, 20, 20}"];
149["container_8 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
104["add_5 (:add) {1, 64, 20, 20}"];
134["mobilenetv20_features_linearbottleneck13_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
75["mobilenetv20_features_linearbottleneck7_conv0_fwd (:conv) {1, 384, 20, 20}"];
90["mobilenetv20_features_linearbottleneck8_relu1_fwd (:relu) {1, 384, 20, 20}"];
154["mobilenetv20_features_linearbottleneck15_conv1_fwd (:conv) {1, 960, 5, 5}"];
151["mobilenetv20_features_linearbottleneck15_conv0_fwd (:conv) {1, 960, 5, 5}"];
58["mobilenetv20_features_linearbottleneck5_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
135["mobilenetv20_features_linearbottleneck13_relu0_fwd (:relu) {1, 576, 10, 10}"];
55["container_1 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
17["mobilenetv20_features_linearbottleneck0_batchnorm1_fwd (:batch_norm) {1, 32, 80, 80}"];
22["mobilenetv20_features_linearbottleneck1_batchnorm0_fwd (:batch_norm) {1, 96, 80, 80}"];
52["mobilenetv20_features_linearbottleneck4_relu1_fwd (:relu) {1, 192, 20, 20}"];
111["mobilenetv20_features_linearbottleneck10_conv2_fwd (:conv) {1, 96, 10, 10}"];
155["mobilenetv20_features_linearbottleneck15_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
87["mobilenetv20_features_linearbottleneck8_relu0_fwd (:relu) {1, 384, 20, 20}"];
21["mobilenetv20_features_linearbottleneck1_conv0_fwd (:conv) {1, 96, 80, 80}"];
96["mobilenetv20_features_linearbottleneck9_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
128["mobilenetv20_features_linearbottleneck12_relu1_fwd (:relu) {1, 576, 10, 10}"];
159["container_9 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
169["mobilenetv20_features_conv1_fwd (:conv) {1, 1280, 5, 5}"];
119["mobilenetv20_features_linearbottleneck11_conv2_fwd (:conv) {1, 96, 10, 10}"];
148["mobilenetv20_features_linearbottleneck14_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
82["mobilenetv20_features_linearbottleneck7_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
36["mobilenetv20_features_linearbottleneck2_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
138["mobilenetv20_features_linearbottleneck13_relu1_fwd (:relu) {1, 576, 5, 5}"];
53["mobilenetv20_features_linearbottleneck4_conv2_fwd (:conv) {1, 32, 20, 20}"];
50["mobilenetv20_features_linearbottleneck4_conv1_fwd (:conv) {1, 192, 20, 20}"];
66["add_2 (:add) {1, 32, 20, 20}"];
24["mobilenetv20_features_linearbottleneck1_conv1_fwd (:conv) {1, 96, 40, 40}"];
10["mobilenetv20_features_conv0_fwd (:conv) {1, 32, 80, 80}"];
150["add_8 (:add) {1, 160, 5, 5}"];
35["mobilenetv20_features_linearbottleneck2_conv2_fwd (:conv) {1, 24, 40, 40}"];
70["mobilenetv20_features_linearbottleneck6_conv1_fwd (:conv) {1, 192, 20, 20}"];
56["add_1 (:add) {1, 32, 20, 20}"];
161["mobilenetv20_features_linearbottleneck16_conv0_fwd (:conv) {1, 960, 5, 5}"];
102["mobilenetv20_features_linearbottleneck9_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
142["mobilenetv20_features_linearbottleneck14_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
62["mobilenetv20_features_linearbottleneck5_relu1_fwd (:relu) {1, 192, 20, 20}"];
105["mobilenetv20_features_linearbottleneck10_conv0_fwd (:conv) {1, 384, 20, 20}"];
127["mobilenetv20_features_linearbottleneck12_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
121["container_6 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
49["mobilenetv20_features_linearbottleneck4_relu0_fwd (:relu) {1, 192, 20, 20}"];
125["mobilenetv20_features_linearbottleneck12_relu0_fwd (:relu) {1, 576, 10, 10}"];
160["add_9 (:add) {1, 160, 5, 5}"];
27["mobilenetv20_features_linearbottleneck1_conv2_fwd (:conv) {1, 24, 40, 40}"];
84["add_3 (:add) {1, 64, 20, 20}"];
94["add_4 (:add) {1, 64, 20, 20}"];
92["mobilenetv20_features_linearbottleneck8_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
42["mobilenetv20_features_linearbottleneck3_conv1_fwd (:conv) {1, 144, 20, 20}"];
19["mobilenetv20_features_linearbottleneck0_conv2_fwd (:conv) {1, 16, 80, 80}"];
67["mobilenetv20_features_linearbottleneck6_conv0_fwd (:conv) {1, 192, 20, 20}"];
93["container_4 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
106["mobilenetv20_features_linearbottleneck10_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
65["container_2 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
101["mobilenetv20_features_linearbottleneck9_conv2_fwd (:conv) {1, 64, 20, 20}"];
156["mobilenetv20_features_linearbottleneck15_relu1_fwd (:relu) {1, 960, 5, 5}"];
129["mobilenetv20_features_linearbottleneck12_conv2_fwd (:conv) {1, 96, 10, 10}"];
114["mobilenetv20_features_linearbottleneck11_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
118["mobilenetv20_features_linearbottleneck11_relu1_fwd (:relu) {1, 576, 10, 10}"];
54["mobilenetv20_features_linearbottleneck4_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
18["mobilenetv20_features_linearbottleneck0_relu1_fwd (:relu) {1, 32, 80, 80}"];
79["mobilenetv20_features_linearbottleneck7_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
136["mobilenetv20_features_linearbottleneck13_conv1_fwd (:conv) {1, 576, 5, 5}"];
172["mobilenetv20_features_pool0_fwd (:global_avg_pool) {1, 1280, 1, 1}"];
115["mobilenetv20_features_linearbottleneck11_relu0_fwd (:relu) {1, 576, 10, 10}"];
107["mobilenetv20_features_linearbottleneck10_relu0_fwd (:relu) {1, 384, 20, 20}"];
164["mobilenetv20_features_linearbottleneck16_conv1_fwd (:conv) {1, 960, 5, 5}"];
16["mobilenetv20_features_linearbottleneck0_conv1_fwd (:conv) {1, 32, 80, 80}"];
173 --> 174;
172 --> 173;
171 --> 172;
170 --> 171;
169 --> 170;
168 --> 169;
167 --> 168;
166 --> 167;
165 --> 166;
164 --> 165;
163 --> 164;
162 --> 163;
161 --> 162;
160 --> 161;
159 --> 160;
150 --> 159;
158 --> 159;
157 --> 158;
156 --> 157;
155 --> 156;
154 --> 155;
153 --> 154;
152 --> 153;
151 --> 152;
150 --> 151;
149 --> 150;
140 --> 149;
148 --> 149;
147 --> 148;
146 --> 147;
145 --> 146;
144 --> 145;
143 --> 144;
142 --> 143;
141 --> 142;
140 --> 141;
139 --> 140;
138 --> 139;
137 --> 138;
136 --> 137;
135 --> 136;
134 --> 135;
133 --> 134;
132 --> 133;
131 --> 132;
122 --> 131;
130 --> 131;
129 --> 130;
128 --> 129;
127 --> 128;
126 --> 127;
125 --> 126;
124 --> 125;
123 --> 124;
122 --> 123;
121 --> 122;
112 --> 121;
120 --> 121;
119 --> 120;
118 --> 119;
117 --> 118;
116 --> 117;
115 --> 116;
114 --> 115;
113 --> 114;
112 --> 113;
111 --> 112;
110 --> 111;
109 --> 110;
108 --> 109;
107 --> 108;
106 --> 107;
105 --> 106;
104 --> 105;
103 --> 104;
94 --> 103;
102 --> 103;
101 --> 102;
100 --> 101;
99 --> 100;
98 --> 99;
97 --> 98;
96 --> 97;
95 --> 96;
94 --> 95;
93 --> 94;
84 --> 93;
92 --> 93;
91 --> 92;
90 --> 91;
89 --> 90;
88 --> 89;
87 --> 88;
86 --> 87;
85 --> 86;
84 --> 85;
83 --> 84;
74 --> 83;
82 --> 83;
81 --> 82;
80 --> 81;
79 --> 80;
78 --> 79;
77 --> 78;
76 --> 77;
75 --> 76;
74 --> 75;
73 --> 74;
72 --> 73;
71 --> 72;
70 --> 71;
69 --> 70;
68 --> 69;
67 --> 68;
66 --> 67;
65 --> 66;
56 --> 65;
64 --> 65;
63 --> 64;
62 --> 63;
61 --> 62;
60 --> 61;
59 --> 60;
58 --> 59;
57 --> 58;
56 --> 57;
55 --> 56;
46 --> 55;
54 --> 55;
53 --> 54;
52 --> 53;
51 --> 52;
50 --> 51;
49 --> 50;
48 --> 49;
47 --> 48;
46 --> 47;
45 --> 46;
44 --> 45;
43 --> 44;
42 --> 43;
41 --> 42;
40 --> 41;
39 --> 40;
38 --> 39;
37 --> 38;
28 --> 37;
36 --> 37;
35 --> 36;
34 --> 35;
33 --> 34;
32 --> 33;
31 --> 32;
30 --> 31;
29 --> 30;
28 --> 29;
27 --> 28;
26 --> 27;
25 --> 26;
24 --> 25;
23 --> 24;
22 --> 23;
21 --> 22;
20 --> 21;
19 --> 20;
18 --> 19;
17 --> 18;
16 --> 17;
15 --> 16;
14 --> 15;
13 --> 14;
12 --> 13;
11 --> 12;
10 --> 11;
9 --> 10;
{_popped, cnn_base} = cnn_base |> Axon.pop_node()
{_popped, cnn_base} = cnn_base |> Axon.pop_node()
{%Axon.Node{
   id: 173,
   name: #Function<67.122028880/2 in Axon.name/2>,
   mode: :both,
   parent: [172],
   parameters: [
     %Axon.Parameter{
       name: "kernel",
       shape: #Function<7.122028880/1 in Axon.conv/3>,
       initializer: #Function<3.99264144/3 in Axon.Initializers.glorot_uniform/1>,
       type: {:f, 32},
       frozen: false
     }
   ],
   args: [:layer, :parameter],
   op: :conv,
   policy: p=f32 c=f32 o=f32,
   hooks: [],
   opts: [
     strides: [1, 1],
     padding: [{0, 0}, {0, 0}],
     input_dilation: 1,
     kernel_dilation: [1, 1],
     feature_group_size: 1,
     channels: :first
   ],
   op_name: :conv,
   stacktrace: [
     {Axon, :layer, 3, [file: 'lib/axon.ex', line: 338]},
     {Axon, :conv, 3, [file: 'lib/axon.ex', line: 912]},
     {AxonOnnx.Deserialize, :recur_nodes, 2, [file: 'lib/axon_onnx/deserialize.ex', line: 961]},
     {Enum, :"-reduce/3-lists^foldl/2-0-", 3, [file: 'lib/enum.ex', line: 2468]},
     {AxonOnnx.Deserialize, :graph_to_axon, 2, [file: 'lib/axon_onnx/deserialize.ex', line: 43]},
     {AxonOnnx.Deserialize, :to_axon, 2, [file: 'lib/axon_onnx/deserialize.ex', line: 26]}
   ]
 },
 #Axon<
   inputs: %{"data" => {1, 3, 224, 224}}
   outputs: "mobilenetv20_features_pool0_fwd"
   nodes: 164
 >}
Axon.Display.as_graph(cnn_base, input_template)
graph TD;
33["mobilenetv20_features_linearbottleneck2_batchnorm1_fwd (:batch_norm) {1, 144, 40, 40}"];
168["mobilenetv20_features_linearbottleneck16_batchnorm2_fwd (:batch_norm) {1, 320, 5, 5}"];
117["mobilenetv20_features_linearbottleneck11_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
12["mobilenetv20_features_relu0_fwd (:relu) {1, 32, 80, 80}"];
157["mobilenetv20_features_linearbottleneck15_conv2_fwd (:conv) {1, 160, 5, 5}"];
132["add_7 (:add) {1, 96, 10, 10}"];
73["mobilenetv20_features_linearbottleneck6_conv2_fwd (:conv) {1, 64, 20, 20}"];
44["mobilenetv20_features_linearbottleneck3_relu1_fwd (:relu) {1, 144, 20, 20}"];
124["mobilenetv20_features_linearbottleneck12_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
170["mobilenetv20_features_batchnorm1_fwd (:batch_norm) {1, 1280, 5, 5}"];
23["mobilenetv20_features_linearbottleneck1_relu0_fwd (:relu) {1, 96, 80, 80}"];
29["mobilenetv20_features_linearbottleneck2_conv0_fwd (:conv) {1, 144, 40, 40}"];
47["mobilenetv20_features_linearbottleneck4_conv0_fwd (:conv) {1, 192, 20, 20}"];
89["mobilenetv20_features_linearbottleneck8_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
61["mobilenetv20_features_linearbottleneck5_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
30["mobilenetv20_features_linearbottleneck2_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
43["mobilenetv20_features_linearbottleneck3_batchnorm1_fwd (:batch_norm) {1, 144, 20, 20}"];
163["mobilenetv20_features_linearbottleneck16_relu0_fwd (:relu) {1, 960, 5, 5}"];
39["mobilenetv20_features_linearbottleneck3_conv0_fwd (:conv) {1, 144, 40, 40}"];
131["container_7 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
45["mobilenetv20_features_linearbottleneck3_conv2_fwd (:conv) {1, 32, 20, 20}"];
48["mobilenetv20_features_linearbottleneck4_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
145["mobilenetv20_features_linearbottleneck14_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
171["mobilenetv20_features_relu1_fwd (:relu) {1, 1280, 5, 5}"];
57["mobilenetv20_features_linearbottleneck5_conv0_fwd (:conv) {1, 192, 20, 20}"];
143["mobilenetv20_features_linearbottleneck14_relu0_fwd (:relu) {1, 960, 5, 5}"];
113["mobilenetv20_features_linearbottleneck11_conv0_fwd (:conv) {1, 576, 10, 10}"];
26["mobilenetv20_features_linearbottleneck1_relu1_fwd (:relu) {1, 96, 40, 40}"];
69["mobilenetv20_features_linearbottleneck6_relu0_fwd (:relu) {1, 192, 20, 20}"];
88["mobilenetv20_features_linearbottleneck8_conv1_fwd (:conv) {1, 384, 20, 20}"];
166["mobilenetv20_features_linearbottleneck16_relu1_fwd (:relu) {1, 960, 5, 5}"];
144["mobilenetv20_features_linearbottleneck14_conv1_fwd (:conv) {1, 960, 5, 5}"];
141["mobilenetv20_features_linearbottleneck14_conv0_fwd (:conv) {1, 960, 5, 5}"];
152["mobilenetv20_features_linearbottleneck15_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
63["mobilenetv20_features_linearbottleneck5_conv2_fwd (:conv) {1, 32, 20, 20}"];
165["mobilenetv20_features_linearbottleneck16_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
100["mobilenetv20_features_linearbottleneck9_relu1_fwd (:relu) {1, 384, 20, 20}"];
71["mobilenetv20_features_linearbottleneck6_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
133["mobilenetv20_features_linearbottleneck13_conv0_fwd (:conv) {1, 576, 10, 10}"];
46["mobilenetv20_features_linearbottleneck3_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
31["mobilenetv20_features_linearbottleneck2_relu0_fwd (:relu) {1, 144, 40, 40}"];
158["mobilenetv20_features_linearbottleneck15_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
167["mobilenetv20_features_linearbottleneck16_conv2_fwd (:conv) {1, 320, 5, 5}"];
98["mobilenetv20_features_linearbottleneck9_conv1_fwd (:conv) {1, 384, 20, 20}"];
81["mobilenetv20_features_linearbottleneck7_conv2_fwd (:conv) {1, 64, 20, 20}"];
11["mobilenetv20_features_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
37["container_0 (:container) {{1, 24, 40, 40}, {1, 24, 40, 40}}"];
153["mobilenetv20_features_linearbottleneck15_relu0_fwd (:relu) {1, 960, 5, 5}"];
9[/"data (:input) {1, 3, 160, 160}"/];
76["mobilenetv20_features_linearbottleneck7_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
32["mobilenetv20_features_linearbottleneck2_conv1_fwd (:conv) {1, 144, 40, 40}"];
34["mobilenetv20_features_linearbottleneck2_relu1_fwd (:relu) {1, 144, 40, 40}"];
25["mobilenetv20_features_linearbottleneck1_batchnorm1_fwd (:batch_norm) {1, 96, 40, 40}"];
28["mobilenetv20_features_linearbottleneck1_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
85["mobilenetv20_features_linearbottleneck8_conv0_fwd (:conv) {1, 384, 20, 20}"];
91["mobilenetv20_features_linearbottleneck8_conv2_fwd (:conv) {1, 64, 20, 20}"];
120["mobilenetv20_features_linearbottleneck11_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
64["mobilenetv20_features_linearbottleneck5_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
146["mobilenetv20_features_linearbottleneck14_relu1_fwd (:relu) {1, 960, 5, 5}"];
109["mobilenetv20_features_linearbottleneck10_batchnorm1_fwd (:batch_norm) {1, 384, 10, 10}"];
72["mobilenetv20_features_linearbottleneck6_relu1_fwd (:relu) {1, 192, 20, 20}"];
110["mobilenetv20_features_linearbottleneck10_relu1_fwd (:relu) {1, 384, 10, 10}"];
68["mobilenetv20_features_linearbottleneck6_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
86["mobilenetv20_features_linearbottleneck8_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
116["mobilenetv20_features_linearbottleneck11_conv1_fwd (:conv) {1, 576, 10, 10}"];
51["mobilenetv20_features_linearbottleneck4_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
80["mobilenetv20_features_linearbottleneck7_relu1_fwd (:relu) {1, 384, 20, 20}"];
38["add_0 (:add) {1, 24, 40, 40}"];
137["mobilenetv20_features_linearbottleneck13_batchnorm1_fwd (:batch_norm) {1, 576, 5, 5}"];
126["mobilenetv20_features_linearbottleneck12_conv1_fwd (:conv) {1, 576, 10, 10}"];
13["mobilenetv20_features_linearbottleneck0_conv0_fwd (:conv) {1, 32, 80, 80}"];
59["mobilenetv20_features_linearbottleneck5_relu0_fwd (:relu) {1, 192, 20, 20}"];
40["mobilenetv20_features_linearbottleneck3_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
123["mobilenetv20_features_linearbottleneck12_conv0_fwd (:conv) {1, 576, 10, 10}"];
77["mobilenetv20_features_linearbottleneck7_relu0_fwd (:relu) {1, 384, 20, 20}"];
95["mobilenetv20_features_linearbottleneck9_conv0_fwd (:conv) {1, 384, 20, 20}"];
130["mobilenetv20_features_linearbottleneck12_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
41["mobilenetv20_features_linearbottleneck3_relu0_fwd (:relu) {1, 144, 40, 40}"];
20["mobilenetv20_features_linearbottleneck0_batchnorm2_fwd (:batch_norm) {1, 16, 80, 80}"];
78["mobilenetv20_features_linearbottleneck7_conv1_fwd (:conv) {1, 384, 20, 20}"];
15["mobilenetv20_features_linearbottleneck0_relu0_fwd (:relu) {1, 32, 80, 80}"];
147["mobilenetv20_features_linearbottleneck14_conv2_fwd (:conv) {1, 160, 5, 5}"];
14["mobilenetv20_features_linearbottleneck0_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
108["mobilenetv20_features_linearbottleneck10_conv1_fwd (:conv) {1, 384, 10, 10}"];
60["mobilenetv20_features_linearbottleneck5_conv1_fwd (:conv) {1, 192, 20, 20}"];
162["mobilenetv20_features_linearbottleneck16_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
103["container_5 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
122["add_6 (:add) {1, 96, 10, 10}"];
140["mobilenetv20_features_linearbottleneck13_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
83["container_3 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
74["mobilenetv20_features_linearbottleneck6_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
139["mobilenetv20_features_linearbottleneck13_conv2_fwd (:conv) {1, 160, 5, 5}"];
99["mobilenetv20_features_linearbottleneck9_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
112["mobilenetv20_features_linearbottleneck10_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
97["mobilenetv20_features_linearbottleneck9_relu0_fwd (:relu) {1, 384, 20, 20}"];
149["container_8 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
104["add_5 (:add) {1, 64, 20, 20}"];
134["mobilenetv20_features_linearbottleneck13_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
75["mobilenetv20_features_linearbottleneck7_conv0_fwd (:conv) {1, 384, 20, 20}"];
90["mobilenetv20_features_linearbottleneck8_relu1_fwd (:relu) {1, 384, 20, 20}"];
154["mobilenetv20_features_linearbottleneck15_conv1_fwd (:conv) {1, 960, 5, 5}"];
151["mobilenetv20_features_linearbottleneck15_conv0_fwd (:conv) {1, 960, 5, 5}"];
58["mobilenetv20_features_linearbottleneck5_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
135["mobilenetv20_features_linearbottleneck13_relu0_fwd (:relu) {1, 576, 10, 10}"];
55["container_1 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
17["mobilenetv20_features_linearbottleneck0_batchnorm1_fwd (:batch_norm) {1, 32, 80, 80}"];
22["mobilenetv20_features_linearbottleneck1_batchnorm0_fwd (:batch_norm) {1, 96, 80, 80}"];
52["mobilenetv20_features_linearbottleneck4_relu1_fwd (:relu) {1, 192, 20, 20}"];
111["mobilenetv20_features_linearbottleneck10_conv2_fwd (:conv) {1, 96, 10, 10}"];
155["mobilenetv20_features_linearbottleneck15_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
87["mobilenetv20_features_linearbottleneck8_relu0_fwd (:relu) {1, 384, 20, 20}"];
21["mobilenetv20_features_linearbottleneck1_conv0_fwd (:conv) {1, 96, 80, 80}"];
96["mobilenetv20_features_linearbottleneck9_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
128["mobilenetv20_features_linearbottleneck12_relu1_fwd (:relu) {1, 576, 10, 10}"];
159["container_9 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
169["mobilenetv20_features_conv1_fwd (:conv) {1, 1280, 5, 5}"];
119["mobilenetv20_features_linearbottleneck11_conv2_fwd (:conv) {1, 96, 10, 10}"];
148["mobilenetv20_features_linearbottleneck14_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
82["mobilenetv20_features_linearbottleneck7_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
36["mobilenetv20_features_linearbottleneck2_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
138["mobilenetv20_features_linearbottleneck13_relu1_fwd (:relu) {1, 576, 5, 5}"];
53["mobilenetv20_features_linearbottleneck4_conv2_fwd (:conv) {1, 32, 20, 20}"];
50["mobilenetv20_features_linearbottleneck4_conv1_fwd (:conv) {1, 192, 20, 20}"];
66["add_2 (:add) {1, 32, 20, 20}"];
24["mobilenetv20_features_linearbottleneck1_conv1_fwd (:conv) {1, 96, 40, 40}"];
10["mobilenetv20_features_conv0_fwd (:conv) {1, 32, 80, 80}"];
150["add_8 (:add) {1, 160, 5, 5}"];
35["mobilenetv20_features_linearbottleneck2_conv2_fwd (:conv) {1, 24, 40, 40}"];
70["mobilenetv20_features_linearbottleneck6_conv1_fwd (:conv) {1, 192, 20, 20}"];
56["add_1 (:add) {1, 32, 20, 20}"];
161["mobilenetv20_features_linearbottleneck16_conv0_fwd (:conv) {1, 960, 5, 5}"];
102["mobilenetv20_features_linearbottleneck9_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
142["mobilenetv20_features_linearbottleneck14_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
62["mobilenetv20_features_linearbottleneck5_relu1_fwd (:relu) {1, 192, 20, 20}"];
105["mobilenetv20_features_linearbottleneck10_conv0_fwd (:conv) {1, 384, 20, 20}"];
127["mobilenetv20_features_linearbottleneck12_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
121["container_6 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
49["mobilenetv20_features_linearbottleneck4_relu0_fwd (:relu) {1, 192, 20, 20}"];
125["mobilenetv20_features_linearbottleneck12_relu0_fwd (:relu) {1, 576, 10, 10}"];
160["add_9 (:add) {1, 160, 5, 5}"];
27["mobilenetv20_features_linearbottleneck1_conv2_fwd (:conv) {1, 24, 40, 40}"];
84["add_3 (:add) {1, 64, 20, 20}"];
94["add_4 (:add) {1, 64, 20, 20}"];
92["mobilenetv20_features_linearbottleneck8_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
42["mobilenetv20_features_linearbottleneck3_conv1_fwd (:conv) {1, 144, 20, 20}"];
19["mobilenetv20_features_linearbottleneck0_conv2_fwd (:conv) {1, 16, 80, 80}"];
67["mobilenetv20_features_linearbottleneck6_conv0_fwd (:conv) {1, 192, 20, 20}"];
93["container_4 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
106["mobilenetv20_features_linearbottleneck10_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
65["container_2 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
101["mobilenetv20_features_linearbottleneck9_conv2_fwd (:conv) {1, 64, 20, 20}"];
156["mobilenetv20_features_linearbottleneck15_relu1_fwd (:relu) {1, 960, 5, 5}"];
129["mobilenetv20_features_linearbottleneck12_conv2_fwd (:conv) {1, 96, 10, 10}"];
114["mobilenetv20_features_linearbottleneck11_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
118["mobilenetv20_features_linearbottleneck11_relu1_fwd (:relu) {1, 576, 10, 10}"];
54["mobilenetv20_features_linearbottleneck4_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
18["mobilenetv20_features_linearbottleneck0_relu1_fwd (:relu) {1, 32, 80, 80}"];
79["mobilenetv20_features_linearbottleneck7_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
136["mobilenetv20_features_linearbottleneck13_conv1_fwd (:conv) {1, 576, 5, 5}"];
172["mobilenetv20_features_pool0_fwd (:global_avg_pool) {1, 1280, 1, 1}"];
115["mobilenetv20_features_linearbottleneck11_relu0_fwd (:relu) {1, 576, 10, 10}"];
107["mobilenetv20_features_linearbottleneck10_relu0_fwd (:relu) {1, 384, 20, 20}"];
164["mobilenetv20_features_linearbottleneck16_conv1_fwd (:conv) {1, 960, 5, 5}"];
16["mobilenetv20_features_linearbottleneck0_conv1_fwd (:conv) {1, 32, 80, 80}"];
171 --> 172;
170 --> 171;
169 --> 170;
168 --> 169;
167 --> 168;
166 --> 167;
165 --> 166;
164 --> 165;
163 --> 164;
162 --> 163;
161 --> 162;
160 --> 161;
159 --> 160;
150 --> 159;
158 --> 159;
157 --> 158;
156 --> 157;
155 --> 156;
154 --> 155;
153 --> 154;
152 --> 153;
151 --> 152;
150 --> 151;
149 --> 150;
140 --> 149;
148 --> 149;
147 --> 148;
146 --> 147;
145 --> 146;
144 --> 145;
143 --> 144;
142 --> 143;
141 --> 142;
140 --> 141;
139 --> 140;
138 --> 139;
137 --> 138;
136 --> 137;
135 --> 136;
134 --> 135;
133 --> 134;
132 --> 133;
131 --> 132;
122 --> 131;
130 --> 131;
129 --> 130;
128 --> 129;
127 --> 128;
126 --> 127;
125 --> 126;
124 --> 125;
123 --> 124;
122 --> 123;
121 --> 122;
112 --> 121;
120 --> 121;
119 --> 120;
118 --> 119;
117 --> 118;
116 --> 117;
115 --> 116;
114 --> 115;
113 --> 114;
112 --> 113;
111 --> 112;
110 --> 111;
109 --> 110;
108 --> 109;
107 --> 108;
106 --> 107;
105 --> 106;
104 --> 105;
103 --> 104;
94 --> 103;
102 --> 103;
101 --> 102;
100 --> 101;
99 --> 100;
98 --> 99;
97 --> 98;
96 --> 97;
95 --> 96;
94 --> 95;
93 --> 94;
84 --> 93;
92 --> 93;
91 --> 92;
90 --> 91;
89 --> 90;
88 --> 89;
87 --> 88;
86 --> 87;
85 --> 86;
84 --> 85;
83 --> 84;
74 --> 83;
82 --> 83;
81 --> 82;
80 --> 81;
79 --> 80;
78 --> 79;
77 --> 78;
76 --> 77;
75 --> 76;
74 --> 75;
73 --> 74;
72 --> 73;
71 --> 72;
70 --> 71;
69 --> 70;
68 --> 69;
67 --> 68;
66 --> 67;
65 --> 66;
56 --> 65;
64 --> 65;
63 --> 64;
62 --> 63;
61 --> 62;
60 --> 61;
59 --> 60;
58 --> 59;
57 --> 58;
56 --> 57;
55 --> 56;
46 --> 55;
54 --> 55;
53 --> 54;
52 --> 53;
51 --> 52;
50 --> 51;
49 --> 50;
48 --> 49;
47 --> 48;
46 --> 47;
45 --> 46;
44 --> 45;
43 --> 44;
42 --> 43;
41 --> 42;
40 --> 41;
39 --> 40;
38 --> 39;
37 --> 38;
28 --> 37;
36 --> 37;
35 --> 36;
34 --> 35;
33 --> 34;
32 --> 33;
31 --> 32;
30 --> 31;
29 --> 30;
28 --> 29;
27 --> 28;
26 --> 27;
25 --> 26;
24 --> 25;
23 --> 24;
22 --> 23;
21 --> 22;
20 --> 21;
19 --> 20;
18 --> 19;
17 --> 18;
16 --> 17;
15 --> 16;
14 --> 15;
13 --> 14;
12 --> 13;
11 --> 12;
10 --> 11;
9 --> 10;
model =
  cnn_base
  |> Axon.namespace("feature_extractor")
  |> Axon.freeze()
  |> Axon.global_avg_pool(channels: :first)
  |> Axon.dense(1)
#Axon<
  inputs: %{"data" => {1, 3, 224, 224}}
  outputs: "dense_0"
  nodes: 167
>
loss = &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2, reduction: :mean, from_logits: true)
optimizer = Axon.Optimizers.adam(1.0e-4)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(train_pipeline, %{"feature_extractor" => cnn_base_params}, epochs: 100)

16:57:29.478 [warn] found unexpected key in the initial parameters map: "mobilenetv20_output_pred_fwd"
eval_model = model |> Axon.sigmoid()

eval_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)
Batch: 23, accuracy: 0.9644719
%{
  0 => %{
    "accuracy" => #Nx.Tensor<
      f32
      EXLA.Backend
      0.9644719362258911
    >
  }
}

Fine Tuning

model = model |> Axon.unfreeze(up: 50)
#Axon<
  inputs: %{"data" => {1, 3, 224, 224}}
  outputs: "dense_0"
  nodes: 166
>
loss = &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2, reduction: :mean, from_logits: true)
optimizer = Axon.Optimizers.rmsprop(1.0e-5)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(train_pipeline, %{"feature_extractor" => cnn_base_params}, epochs: 100)

13:11:18.803 [warn] found unexpected key in the initial parameters map: "mobilenetv20_output_pred_fwd"

13:11:23.337 [info] ptxas warning : Registers are spilled to local memory in function 'fusion_913', 100 bytes spill stores, 100 bytes spill loads

Epoch: 0, Batch: 700, accuracy: 0.8996100 loss: 0.2287182
Batch: 23, accuracy: 0.9596355 loss: 0.0857394
Epoch: 1, Batch: 700, accuracy: 0.9407119 loss: 0.1841684
Batch: 23, accuracy: 0.9661460 loss: 0.0687413
Epoch: 2, Batch: 700, accuracy: 0.9509653 loss: 0.1625399
Batch: 23, accuracy: 0.9687502 loss: 0.0692015
Epoch: 3, Batch: 700, accuracy: 0.9566265 loss: 0.1480721
Batch: 23, accuracy: 0.9752605 loss: 0.0614658
Epoch: 4, Batch: 700, accuracy: 0.9619762 loss: 0.1373688
Batch: 23, accuracy: 0.9739585 loss: 0.0710718
Epoch: 5, Batch: 700, accuracy: 0.9660777 loss: 0.1285189
Batch: 23, accuracy: 0.9739584 loss: 0.0710575
Epoch: 6, Batch: 700, accuracy: 0.9713383 loss: 0.1209094
Batch: 23, accuracy: 0.9765627 loss: 0.0717105
Epoch: 7, Batch: 700, accuracy: 0.9726299 loss: 0.1142527
Batch: 23, accuracy: 0.9778647 loss: 0.0800760
Epoch: 8, Batch: 700, accuracy: 0.9770432 loss: 0.1082386
Batch: 23, accuracy: 0.9752605 loss: 0.0922253
Epoch: 9, Batch: 700, accuracy: 0.9777576 loss: 0.1030978
Batch: 23, accuracy: 0.9709822 loss: 0.1120541
%{
  "dense_0" => %{
    "bias" => #Nx.Tensor<
      f32[1]
      EXLA.Backend
      [-0.0010021841153502464]
    >,
    "kernel" => #Nx.Tensor<
      f32[1280][1]
      EXLA.Backend
      [
        [0.003931004554033279],
        [0.05454326793551445],
        [-0.03750719130039215],
        [0.00593008566647768],
        [-0.0538177452981472],
        [-0.05506449565291405],
        [0.047323524951934814],
        [0.038613103330135345],
        [0.011689354665577412],
        [0.002227860502898693],
        [0.04574345424771309],
        [-0.008441291749477386],
        [0.054752327501773834],
        [-0.009868746623396873],
        [-0.036137595772743225],
        [0.023157991468906403],
        [-0.016265619546175003],
        [0.0349438339471817],
        [0.055996350944042206],
        [0.05815635994076729],
        [0.05296924710273743],
        [-0.04407293722033501],
        [0.013356780633330345],
        [-0.012349767610430717],
        [0.03428534418344498],
        [-0.04150447994470596],
        [0.01897531934082508],
        [0.05189206078648567],
        [0.03463485836982727],
        [-0.012466801330447197],
        [-0.011061359196901321],
        [0.047283440828323364],
        [-0.0013741496950387955],
        [-0.0390496663749218],
        [-0.05481628701090813],
        [0.04738018289208412],
        [0.02277243323624134],
        [-0.029113439843058586],
        [0.019017575308680534],
        [0.02643076330423355],
        [0.037657223641872406],
        [-0.06337488442659378],
        [-0.006380930542945862],
        [0.0020927051082253456],
        [0.015787983313202858],
        [-0.0422554537653923],
        [-0.04168154299259186],
        ...
      ]
    >
  },
  "feature_extractor" => %{
    "mobilenetv20_features_linearbottleneck14_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-0.05331932008266449, -0.09149160981178284, -0.1113930493593216, 0.024654751643538475, -0.018937254324555397, 0.03685024753212929, -0.05549876391887665, 0.055007148534059525, -0.03068343736231327, 0.009129343554377556, 0.14446669816970825, 0.028845686465501785, -0.047605618834495544, -0.021105285733938217, -0.012838606722652912, 0.03015807829797268, 0.02690724842250347, -0.05568072944879532, 0.027742793783545494, -0.05115864425897598, -0.016480594873428345, -0.032794319093227386, 0.25051823258399963, -0.1163983941078186, -0.1583947092294693, -0.012538143433630466, 0.13880954682826996, 0.053955089300870895, 0.08988376706838608, -0.08362587541341782, 0.06670179218053818, -0.18035581707954407, -0.053687501698732376, -0.05454595759510994, -0.033035676926374435, 0.0295602697879076, -0.009064313024282455, -0.05049778148531914, -0.05471234768629074, -0.051531895995140076, 0.04432453215122223, -0.0075772786512970924, -0.009681222029030323, 0.24422961473464966, -0.013667000457644463, 0.034691546112298965, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.2059236317873001, 0.22103242576122284, 0.19723041355609894, 0.15672855079174042, 0.1756642609834671, 0.15570832788944244, 0.20988349616527557, 0.16357900202274323, 0.17750883102416992, 0.1687183976173401, 0.1295347809791565, 0.11767393350601196, 0.21056729555130005, 0.15419363975524902, 0.18961186707019806, 0.16905507445335388, 0.18616241216659546, 0.2066846638917923, 0.16887988150119781, 0.17675012350082397, 0.15337537229061127, 0.18424303829669952, 0.09231094270944595, 0.23877179622650146, 0.1938796490430832, 0.15088878571987152, 0.13759241998195648, 0.145432248711586, 0.12895819544792175, 0.22074861824512482, 0.1405794322490692, 0.4638647437095642, 0.2596072256565094, 0.2193414717912674, 0.18216955661773682, 0.16033369302749634, 0.18403029441833496, 0.21628527343273163, 0.23245695233345032, 0.23695005476474762, 0.13848741352558136, 0.1762000471353531, 0.1793013960123062, 0.15583865344524384, 0.2042112797498703, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-0.0023206884507089853, -0.004551870748400688, 0.0034389428328722715, 0.0015984668862074614, -0.002293417463079095, -0.03154199197888374, 0.015855349600315094, -0.021844547241926193, -0.0030182870104908943, -0.008290023542940617, -0.008337621577084064, 0.003635477740317583, -0.02967023104429245, 0.009609879925847054, -0.00535233598202467, -0.03133776783943176, -0.034155990928411484, -0.0017219285946339369, -0.02915215492248535, -0.0034448853693902493, -0.0022068622056394815, -0.003281062701717019, 0.003420598339289427, 0.00949984323233366, 8.879394154064357e-4, 0.002022636355832219, -0.02363204024732113, -0.017934758216142654, -0.02237909846007824, -0.05140375345945358, 4.167649312876165e-4, -0.006942620035260916, 0.011951959691941738, -0.0031885183416306973, -9.77475312538445e-4, -0.0014110811753198504, -0.03238283097743988, -0.005639346782118082, -0.0070093655958771706, -0.005301866214722395, 0.0020262207835912704, -0.03578586131334305, -0.0526641346514225, -0.009258503094315529, ...]
      >,
      "var" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [4.962424281984568e-4, 2.4244154337793589e-4, 2.3794326989445835e-4, 2.723340585362166e-4, 3.9763047243468463e-4, 5.111433565616608e-4, 4.7444249503314495e-4, 3.682477108668536e-4, 4.1800420149229467e-4, 7.685577729716897e-4, 3.6789028672501445e-4, 2.2397300926968455e-4, 4.4611067278310657e-4, 2.4524933542124927e-4, 0.0011620114091783762, 5.935822846367955e-4, 7.234420045278966e-4, 3.640245122369379e-4, 5.758852930739522e-4, 3.4469107049517334e-4, 2.2663685376755893e-4, 4.212635103613138e-4, 2.179306757170707e-4, 4.0046244976110756e-4, 1.659863191889599e-4, 8.642849934403785e-6, 7.320880540646613e-4, 3.3074370003305376e-4, 4.445896774996072e-4, 7.599603850394487e-4, 3.4124808735214174e-4, 0.0010172502370551229, 7.496927282772958e-4, 4.7150562750175595e-4, 7.05038895830512e-4, 2.509213227313012e-4, 5.495413788594306e-4, 4.112518217880279e-4, 0.0011761350324377418, 6.135390140116215e-4, 4.0227093268185854e-4, 6.605967064388096e-4, 5.76652237214148e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck15_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[160][960][1][1]
        EXLA.Backend
        [
          [
            [
              [0.0034362361766397953]
            ],
            [
              [-0.09938374906778336]
            ],
            [
              [0.0020771357230842113]
            ],
            [
              [0.1406560242176056]
            ],
            [
              [0.03093625232577324]
            ],
            [
              [-0.05371454730629921]
            ],
            [
              [-0.03501325473189354]
            ],
            [
              [-0.01952589489519596]
            ],
            [
              [0.04774273931980133]
            ],
            [
              [-0.011959793046116829]
            ],
            [
              [0.021270673722028732]
            ],
            [
              [-0.04958450421690941]
            ],
            [
              [0.018290305510163307]
            ],
            [
              [0.01978115551173687]
            ],
            [
              [-0.01479385793209076]
            ],
            [
              [0.026327388361096382]
            ],
            [
              [-0.02063232660293579]
            ],
            [
              [-0.032455939799547195]
            ],
            [
              [-0.0023264160845428705]
            ],
            [
              [-0.027306411415338516]
            ],
            [
              [0.012356900610029697]
            ],
            [
              [0.008741793222725391]
            ],
            [
              [-0.05494425818324089]
            ],
            [
              [-0.007671859581023455]
            ],
            [
              [-0.03244822099804878]
            ],
            [
              [-0.02167103998363018]
            ],
            [
              [0.12303319573402405]
            ],
            [
              [0.05844210460782051]
            ],
            [
              [-0.011791139841079712]
            ],
            [
              [-0.03864377364516258]
            ],
            [
              [-0.059811562299728394]
            ],
            [
              [0.03294980898499489]
            ],
            [
              [-0.01898467168211937]
            ],
            [
              [-0.044919755309820175]
            ],
            [
              [-0.006398592609912157]
            ],
            [
              [-0.052493866533041]
            ],
            [
              [-0.013703064993023872]
            ],
            [
              [-0.03498911112546921]
            ],
            [
              [-0.0736972987651825]
            ],
            [
              [-0.027700163424015045]
            ],
            [
              [-0.04219427332282066]
            ],
            [
              [0.044594526290893555]
            ],
            [
              [0.02880723401904106]
            ],
            [
              [-0.03869602456688881]
            ],
            [
              [0.0056617711670696735]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck12_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][576][1][1]
        EXLA.Backend
        [
          [
            [
              [0.0363660603761673]
            ],
            [
              [0.024831045418977737]
            ],
            [
              [0.0479116328060627]
            ],
            [
              [0.12827005982398987]
            ],
            [
              [-0.0443074107170105]
            ],
            [
              [-0.06308425217866898]
            ],
            [
              [-0.02774641662836075]
            ],
            [
              [-0.06254592537879944]
            ],
            [
              [0.10319916903972626]
            ],
            [
              [0.043877702206373215]
            ],
            [
              [0.05680917948484421]
            ],
            [
              [0.02603430487215519]
            ],
            [
              [-0.0064709968864917755]
            ],
            [
              [-0.03997242823243141]
            ],
            [
              [0.0566127672791481]
            ],
            [
              [-0.06474670022726059]
            ],
            [
              [0.005571999587118626]
            ],
            [
              [-0.004118780605494976]
            ],
            [
              [-0.052782006561756134]
            ],
            [
              [0.027868814766407013]
            ],
            [
              [0.006087758112698793]
            ],
            [
              [-0.07501402497291565]
            ],
            [
              [-0.07516537606716156]
            ],
            [
              [0.03123542107641697]
            ],
            [
              [-0.06451807916164398]
            ],
            [
              [-0.036801498383283615]
            ],
            [
              [-0.047034330666065216]
            ],
            [
              [0.011775856837630272]
            ],
            [
              [-0.008211801759898663]
            ],
            [
              [0.029848778620362282]
            ],
            [
              [0.019432151690125465]
            ],
            [
              [0.07325898110866547]
            ],
            [
              [0.017153287306427956]
            ],
            [
              [0.0742131844162941]
            ],
            [
              [-0.04829411208629608]
            ],
            [
              [-0.07461646944284439]
            ],
            [
              [0.014501615427434444]
            ],
            [
              [-0.050393667072057724]
            ],
            [
              [-0.013460940681397915]
            ],
            [
              [0.08050172030925751]
            ],
            [
              [0.0014683185145258904]
            ],
            [
              [-0.04553446173667908]
            ],
            [
              [-0.008874980732798576]
            ],
            [
              [0.049730490893125534]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck14_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.11817944794893265, 0.1184530109167099, -0.02600397914648056, 0.06703782826662064, 0.06181669607758522, 0.08450127393007278, -0.0031556233298033476, 0.0670972689986229, 0.07826171815395355, 0.013668162748217583, 0.035489194095134735, -0.04560278728604317, 0.09393332898616791, -0.01878989301621914, 0.12683413922786713, 0.0860162153840065, 0.09113430231809616, 0.09262648969888687, 0.10391075164079666, 0.0062261237762868404, 0.017815779894590378, 0.07444161921739578, -0.038155876100063324, 0.01571900211274624, 0.01851974055171013, -0.012220378965139389, 0.016366537660360336, 0.02445855736732483, 0.0402345210313797, 0.14945313334465027, 0.024916628375649452, -0.05475921183824539, -0.06402751803398132, 0.10481038689613342, 0.09596998244524002, 0.11351870000362396, 0.08836444467306137, 0.014713205397129059, 0.11521978676319122, 0.1276046633720398, 0.07535979896783829, 0.10306184738874435, 0.13939012587070465, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.08684154599905014, 0.06107446178793907, 0.11373654007911682, 0.10464422404766083, 0.12101417034864426, 0.11171407997608185, 0.12980608642101288, 0.12206578999757767, 0.09823109954595566, 0.13756020367145538, 0.11793699115514755, 0.12215695530176163, 0.11066754907369614, 0.11358427256345749, 0.11370634287595749, 0.11825334280729294, 0.13193202018737793, 0.0866175964474678, 0.13185307383537292, 0.12279760092496872, 0.11785532534122467, 0.10195711255073547, 0.1281587928533554, 0.10871037095785141, 0.16346554458141327, 0.040990568697452545, 0.16911770403385162, 0.1387839913368225, 0.13536106050014496, 0.10003330558538437, 0.12638407945632935, 0.17041975259780884, 0.16789720952510834, 0.08531254529953003, 0.11380861699581146, 0.07166118919849396, 0.13404858112335205, 0.12250235676765442, 0.10781623423099518, 0.07945836335420609, 0.11125458031892776, 0.12861424684524536, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-1.49989478614998e-8, -6.598725121875759e-8, -3.3777659780298563e-8, 6.155921328243608e-10, 1.6680649084221955e-9, -6.975810151743644e-8, -1.0280804474405159e-7, 1.1829909851712728e-7, -1.391223776181505e-7, 1.951114825260447e-8, -3.154044847519799e-8, -7.882130859115932e-8, 7.139960445101678e-8, -6.877448299746902e-8, -1.5163976740950602e-7, 4.35149836164328e-8, -3.990682628796094e-8, 3.889567778969649e-8, 8.67010143679181e-8, 1.5913511219878274e-7, -1.216080569577116e-8, -3.217380850628615e-8, -2.502063409792754e-7, -1.2450051656287542e-8, -8.94942004947552e-8, -1.0849137055402025e-7, -7.23792865642281e-8, -2.6819417797696588e-9, -4.9953573721950306e-8, -8.798366479823017e-8, 4.691241528576029e-8, 1.7972880073102715e-8, -1.1659428622579071e-7, -1.6043891548633837e-8, 1.5714466883309797e-7, 1.3846912061410421e-8, -8.361175218851713e-8, -2.2993779325020114e-8, -2.5269557557550115e-9, 2.5268136027989385e-8, -5.639618905206589e-8, ...]
      >,
      "var" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.015375461429357529, 0.008966661989688873, 0.014436393044888973, 0.018274905160069466, 0.022764142602682114, 0.017717313021421432, 0.028602011501789093, 0.02376588247716427, 0.013949170708656311, 0.014504922553896904, 0.015037551522254944, 0.017101136967539787, 0.014324591495096684, 0.024719757959246635, 0.022843269631266594, 0.016248025000095367, 0.02265304885804653, 0.013783901929855347, 0.020945660769939423, 0.010627896524965763, 0.02857845090329647, 0.019863925874233246, 0.010918163694441319, 0.01824241690337658, 0.061989493668079376, 0.013930430635809898, 0.03650098294019699, 0.01744050718843937, 0.015627671033143997, 0.011513426899909973, 0.020515792071819305, 0.024865739047527313, 0.028932511806488037, 0.013199539855122566, 0.01918703131377697, 0.008621279150247574, 0.022145573049783707, 0.017280828207731247, 0.022399913519620895, 0.01449136808514595, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck16_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.04983936622738838, 0.07602783292531967, 0.051380790770053864, 0.3098520040512085, 0.06193145737051964, 0.1989825963973999, 0.059103287756443024, 0.057949237525463104, 0.08095505088567734, 0.06015932559967041, 0.07324974983930588, 0.21073018014431, 0.06095729023218155, 0.01063812430948019, 0.15157859027385712, 0.05909590423107147, 0.053097400814294815, 0.06868501007556915, 0.05966099351644516, 0.05835151672363281, 0.26296988129615784, 0.06431972235441208, 0.06503307819366455, 0.07718635350465775, 0.0668996199965477, 0.06021757796406746, 0.16841691732406616, 0.11744103580713272, 0.062217265367507935, 0.06440437585115433, 0.06086958199739456, 0.06575987488031387, 0.05889017879962921, -0.12167530506849289, 0.06327949464321136, 0.07213419675827026, 0.061806030571460724, 0.06377823650836945, 0.12375108897686005, 0.04857545346021652, 0.06471194326877594, 0.058631494641304016, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.23891015350818634, 0.23229506611824036, 0.2370145469903946, 0.09444864094257355, 0.22579453885555267, 0.183750182390213, 0.25486111640930176, 0.2427559792995453, 0.14071176946163177, 0.22152307629585266, 0.2263786941766739, 0.15101081132888794, 0.2631012201309204, 0.22168727219104767, 0.1091405600309372, 0.18144215643405914, 0.2794880270957947, 0.239920973777771, 0.2652829885482788, 0.235841765999794, 0.09454084187746048, 0.2351168990135193, 0.15646567940711975, 0.14339980483055115, 0.21100369095802307, 0.24368858337402344, 0.1311427801847458, 0.12481340765953064, 0.2191619724035263, 0.21206896007061005, 0.19799712300300598, 0.18251578509807587, 0.1878426969051361, 0.2655341923236847, 0.22601383924484253, 0.21993128955364227, 0.2297792136669159, 0.2173818200826645, 0.12258932739496231, 0.1422921121120453, 0.22551919519901276, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 960]
        EXLA.Backend
        [4.214608925394714e-4, 8.306027739308774e-4, 4.0714265196584165e-4, -0.0070646279491484165, 6.072276155464351e-4, -0.003040345385670662, 3.0289063579402864e-4, 1.9496645836625248e-4, -0.009553076699376106, 5.691817495971918e-4, 6.218577036634088e-4, -0.002603895030915737, 6.721426034346223e-4, -0.04986010864377022, -0.010595223866403103, 4.088203713763505e-4, 2.3354109725914896e-4, 7.168585434556007e-4, 4.6488005318678916e-4, 9.686480043455958e-4, -0.013233187608420849, 3.088755183853209e-4, -0.0075218770653009415, -0.009166075848042965, 9.288735454902053e-4, 5.843568942509592e-4, -0.01027306541800499, -0.006692919414490461, 5.131419748067856e-4, 6.182753713801503e-4, 6.51774404104799e-4, 0.0010231614578515291, 7.825290085747838e-4, -0.007864231243729591, 4.884237423539162e-4, 6.116394069977105e-4, 4.0676729986444116e-4, 5.098831607028842e-4, -0.013810231350362301, -0.016042225062847137, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 960]
        EXLA.Backend
        [8.36657534364349e-7, 1.6661938389006536e-6, 1.1219271982554346e-6, 4.39523282693699e-5, 1.7165073131764075e-6, 9.104181117436383e-6, 5.49723836229532e-7, 3.117878009106789e-7, 8.054874342633411e-5, 1.3060824812782812e-6, 1.123108063438849e-6, 9.358960596728139e-6, 1.7289175957557745e-6, 6.426551844924688e-4, 8.784420788288116e-5, 6.745207770109118e-7, 5.311425184117979e-7, 1.4183875691742287e-6, 1.3538910934585147e-6, 3.6514857129077427e-6, 9.821475396165624e-5, 5.143819521435944e-7, 3.44959698850289e-5, 5.321364005794749e-5, 2.272704250572133e-6, 1.7950275150724337e-6, 6.734966882504523e-5, 7.387526420643553e-5, 9.372954536956968e-7, 1.3867668258171761e-6, 1.0380003914178815e-6, 1.9461390365904663e-6, 1.6668218449922279e-6, 2.1192214626353234e-4, 1.0493308764125686e-6, 9.715754458738957e-7, 1.0053443020296982e-6, 1.0418542615298065e-6, 1.5705001715105027e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck4_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [-0.16025467216968536, -0.08318334817886353, -0.013605363667011261, 0.016528524458408356, -0.0414351262152195, -0.08624469488859177, 0.08369371294975281, 0.08077652007341385, -0.3613238036632538, 0.0863688737154007, -0.07133538275957108, 0.17072106897830963, -0.006156732328236103, 0.04610714688897133, 0.26338741183280945, 0.25080469250679016, 0.0866481363773346, 0.029331553727388382, 0.13819074630737305, -0.06575538963079453, -5.868459120392799e-4, -0.02378389984369278, -0.03555862605571747, -0.022012412548065186, 0.07437903434038162, 0.07370810955762863, 0.30864641070365906, -0.006008080206811428, 0.15304052829742432, -0.00992464367300272, 0.19627465307712555, 0.07678351551294327, 0.10757749527692795, -0.02487078309059143, -0.040003497153520584, -0.07770790904760361, 0.2532285451889038, 0.16046945750713348, 0.05568472295999527, 0.15549562871456146, -0.038443390280008316, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [0.12059392035007477, 0.18736249208450317, 0.1590438038110733, 0.2091878354549408, 0.1450766772031784, 0.15702469646930695, 0.3902188539505005, 0.25355497002601624, 0.3910222351551056, 0.26179036498069763, 0.189780130982399, 0.14822576940059662, 0.210128515958786, 0.1974102109670639, 0.0950021743774414, 0.1603325754404068, 0.12979912757873535, 0.08638784289360046, 0.11244207620620728, 0.1680508255958557, 0.243540421128273, 0.20751050114631653, 0.1465093046426773, 0.15335243940353394, 0.2690037190914154, 0.1203547790646553, 0.0906066969037056, 0.23154661059379578, 0.1154584214091301, 0.1968846470117569, 0.11754576861858368, 0.1189347431063652, 0.3180309534072876, 0.20275117456912994, 0.1695779711008072, 0.2201462984085083, 0.10936479270458221, 0.0758633241057396, 0.12570320069789886, 0.09598797559738159, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [0.005834498442709446, 0.005320976488292217, 0.011644497513771057, -0.0033567524515092373, 0.006512106861919165, 0.0075160483829677105, -0.011912024579942226, -0.021598124876618385, -0.15417402982711792, -0.001387530704960227, -0.0037699746899306774, 0.01811889372766018, -9.935529669746757e-4, -0.00730853620916605, -0.004980568308383226, 5.785240791738033e-4, 0.0198814794421196, 0.0036787642166018486, -0.018057825043797493, -2.4293076421599835e-4, 0.002015749691054225, 0.013546604663133621, 0.007414540275931358, 1.2083808542229235e-4, -0.003124411916360259, -0.019600532948970795, -0.005146954674273729, -0.002622810425236821, -0.016293222084641457, 0.023455090820789337, -0.00871172547340393, -0.008095859549939632, -0.018160350620746613, -5.122238071635365e-4, 0.0051345909014344215, 0.002093162387609482, -0.01451784186065197, 0.0019358042627573013, 0.0016494555165991187, ...]
      >,
      "var" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [1.82411604328081e-4, 2.6851065922528505e-4, 2.388732973486185e-4, 8.272912818938494e-4, 1.7400244541931897e-4, 2.264920767629519e-4, 0.005600188858807087, 0.0019779913127422333, 0.004019826650619507, 0.002996259368956089, 2.9179989360272884e-4, 9.683401440270245e-4, 6.0461163229774684e-5, 1.024205848807469e-4, 3.014750254806131e-4, 4.56174939245102e-6, 8.186800550902262e-5, 4.768615690409206e-5, 3.702163230627775e-4, 3.5265996120870113e-4, 0.001435167738236487, 4.292532103136182e-4, 1.5066159539856017e-4, 1.9934076408389956e-4, 0.0030081914737820625, 9.441975998925045e-5, 3.100578032899648e-4, 2.9591855127364397e-4, 4.2446120642125607e-4, 8.81795072928071e-4, 1.9674932991620153e-4, 2.7562043396756053e-4, 0.0034803457092493773, 6.995198782533407e-4, 3.369579208083451e-4, 6.203856319189072e-4, 4.934457829222083e-4, 4.415657167555764e-5, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck13_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.21549595892429352, 0.1323530673980713, -0.16010034084320068],
              [0.19547533988952637, -0.11036711931228638, 0.026170458644628525],
              [-0.10824786126613617, 0.0671815574169159, -0.7960876822471619]
            ]
          ],
          [
            [
              [-1.1282756328582764, 0.10512088239192963, 0.05058681219816208],
              [-0.1645849198102951, 0.18541687726974487, -0.3736736476421356],
              [0.18559065461158752, 0.037046074867248535, 0.30303117632865906]
            ]
          ],
          [
            [
              [-0.1953502595424652, 0.21459802985191345, 0.016381336376070976],
              [0.1539907604455948, 0.0909615159034729, 0.0656626969575882],
              [-0.7163212895393372, -0.043639495968818665, 0.44717922806739807]
            ]
          ],
          [
            [
              [-0.12282669544219971, 0.16306449472904205, 0.35451921820640564],
              [0.5785645842552185, 0.035986918956041336, 0.5223076343536377],
              [0.2113819420337677, 0.6414011716842651, 0.3652750849723816]
            ]
          ],
          [
            [
              [0.6755218505859375, -0.0564180389046669, 0.049403801560401917],
              [0.060685895383358, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck1_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[24][96][1][1]
        EXLA.Backend
        [
          [
            [
              [0.09949389100074768]
            ],
            [
              [-0.18654611706733704]
            ],
            [
              [0.13225184381008148]
            ],
            [
              [0.03560471162199974]
            ],
            [
              [-0.020650217309594154]
            ],
            [
              [-0.06614962220191956]
            ],
            [
              [-0.12619785964488983]
            ],
            [
              [-0.17344357073307037]
            ],
            [
              [0.13648797571659088]
            ],
            [
              [0.05030941590666771]
            ],
            [
              [0.0744558572769165]
            ],
            [
              [0.02734171226620674]
            ],
            [
              [0.1788010597229004]
            ],
            [
              [-0.0010324945906177163]
            ],
            [
              [-0.058398596942424774]
            ],
            [
              [-0.30487778782844543]
            ],
            [
              [0.23805657029151917]
            ],
            [
              [-0.04994773864746094]
            ],
            [
              [-0.04996243491768837]
            ],
            [
              [0.049236271530389786]
            ],
            [
              [0.15690600872039795]
            ],
            [
              [-0.085136778652668]
            ],
            [
              [-0.06452983617782593]
            ],
            [
              [-0.03775155171751976]
            ],
            [
              [-0.2167663425207138]
            ],
            [
              [0.20706818997859955]
            ],
            [
              [0.23726293444633484]
            ],
            [
              [-0.1283334344625473]
            ],
            [
              [-0.22454236447811127]
            ],
            [
              [-0.18676628172397614]
            ],
            [
              [0.09907364100217819]
            ],
            [
              [0.22619308531284332]
            ],
            [
              [-0.029751716181635857]
            ],
            [
              [0.16340412199497223]
            ],
            [
              [0.01070564053952694]
            ],
            [
              [0.3250269889831543]
            ],
            [
              [0.21018050611019135]
            ],
            [
              [-0.062177326530218124]
            ],
            [
              [0.12571603059768677]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck6_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [0.3632318675518036, -0.023602070286870003, -0.0614071749150753, 0.5761924386024475, 0.1690974086523056, 0.19036972522735596, 0.0031085966620594263, -0.018681762740015984, 0.038269754499197006, 0.03641441464424133, 0.029138026759028435, -0.016922323033213615, 0.055098824203014374, 0.21938563883304596, 0.22719338536262512, 0.20525230467319489, 0.20284974575042725, 0.07362302392721176, 0.1064985990524292, 0.40782883763313293, -0.00754354614764452, 0.009455508552491665, 0.4253721535205841, 0.6265411972999573, 0.22528603672981262, -0.05191487818956375, 0.6341891288757324, 0.347513884305954, 0.04142335429787636, 0.08137086033821106, 0.35836175084114075, 0.038170408457517624, 0.14619027078151703, 0.0850687101483345, 0.36599797010421753, 0.3969757854938507, 0.09776335954666138, 0.02894003875553608, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[192]
        EXLA.Backend
        [0.1283191293478012, 0.2977066934108734, 0.2804311215877533, 0.20587405562400818, 0.2696685492992401, 0.12131325900554657, 0.30051329731941223, 0.36318743228912354, 0.23120905458927155, 0.22658991813659668, 0.24182915687561035, 0.30583131313323975, 0.20871473848819733, 0.2518334686756134, 0.22621870040893555, 0.20286978781223297, 0.19829389452934265, 0.16485567390918732, 0.18886375427246094, 0.1450076848268509, 0.312345415353775, 0.2126840353012085, 0.2775242030620575, 0.1377149373292923, 0.15688155591487885, 0.3647836446762085, 0.24961578845977783, 0.18082654476165771, 0.17116306722164154, 0.23007430136203766, 0.1871042400598526, 0.2649954557418823, 0.306458443403244, 0.21913877129554749, 0.16099828481674194, 0.1434336155653, 0.12895677983760834, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 192]
        EXLA.Backend
        [-0.045330289751291275, 0.018183870241045952, 0.006856370717287064, 0.06764873117208481, -0.132804274559021, -0.038608551025390625, -0.16889694333076477, 0.07737334817647934, -0.0899917483329773, 0.006626186426728964, 0.026949428021907806, -0.01451886910945177, 0.03275036811828613, -0.08090642094612122, -0.06808585673570633, 0.04256214201450348, 0.043091513216495514, 0.024077443405985832, -0.05164222791790962, -0.03613243252038956, 0.03436368703842163, 0.019569337368011475, 0.012944743037223816, -0.035043563693761826, -0.04563178867101669, -0.028518084436655045, 0.10354693979024887, -0.06964914500713348, 0.05185485631227493, 0.043715570122003555, 0.03401125967502594, -0.14536544680595398, 0.02164241299033165, 0.051821183413267136, 0.06498952955007553, -0.008307303301990032, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 192]
        EXLA.Backend
        [0.0023164052981883287, 0.001152294920757413, 0.0013173851184546947, 0.010171322152018547, 0.012587573379278183, 0.0011759239714592695, 0.0028448356315493584, 0.0010036694584414363, 4.9413344822824e-4, 0.002985553815960884, 0.001194263226352632, 0.0026788776740431786, 0.0015006095636636019, 0.009730218909680843, 0.005227346438914537, 0.009902280755341053, 0.004874127451330423, 0.001359938527457416, 0.001311283791437745, 0.001999074127525091, 0.001209039706736803, 0.0014069834724068642, 0.0021509360522031784, 0.0015126715879887342, 9.573991410434246e-4, 9.09177353605628e-4, 0.014325403608381748, 7.886392413638532e-4, 0.0021764112170785666, 0.0033195368014276028, 0.001884391182102263, 0.0053754388354718685, 0.003156874096021056, 0.0030100925359874964, 0.0011924542486667633, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck1_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-5.599617361440323e-5, -1.388202654197812e-4, 1.4964351430535316e-4, -4.48540726210922e-4, -2.7302927264827304e-5, 1.428744290024042e-4, -1.0991278395522386e-4, 1.3559161743614823e-4, 1.4486558211501688e-4, -5.513934302143753e-4, -9.67518353718333e-5, -2.075896190945059e-4, -1.0409099922981113e-4, -1.4599738642573357e-4, -1.6489351401105523e-4, -9.032035450218245e-5, -1.591532345628366e-4, -1.570975873619318e-4, -7.89462574175559e-5, 3.225041655241512e-5, 8.031473407754675e-5, -3.937935980502516e-5, 2.2796950361225754e-5, -3.125887888018042e-4]
      >,
      "gamma" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.6092047691345215, 0.47812265157699585, 0.693821370601654, 0.45786911249160767, 0.5387675166130066, 0.8261818289756775, 0.5807712078094482, 0.6510066986083984, 0.45650821924209595, 0.7799047827720642, 0.6193627715110779, 0.5857614874839783, 0.4286927878856659, 0.6337913870811462, 0.6911339163780212, 0.4652976393699646, 0.5756664872169495, 0.46622365713119507, 0.5369606018066406, 0.5942122340202332, 0.7727246880531311, 0.6141138076782227, 0.46364495158195496, 0.6175265908241272]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 24]
        EXLA.Backend
        [0.22155123949050903, -0.6191795468330383, -0.5176143646240234, 0.17332279682159424, -0.8719085454940796, -0.39355045557022095, 0.38133496046066284, 1.2457244396209717, -0.21562506258487701, 0.8583738207817078, 0.08692432940006256, -1.6506699323654175, 0.5147061944007874, -0.1741609126329422, 0.9139919281005859, -0.13843558728694916, 0.4120836555957794, 0.21832101047039032, -1.1366162300109863, -0.07164992392063141, -0.516703724861145, -1.7801650762557983, -0.5468620657920837, -0.08122311532497406]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 24]
        EXLA.Backend
        [0.1357666403055191, 0.09551053494215012, 0.19151824712753296, 0.073853999376297, 0.08839397877454758, 0.2828894257545471, 0.09247986972332001, 0.1987546682357788, 0.08774235099554062, 0.2564696669578552, 0.13650986552238464, 0.1383553445339203, 0.06085009127855301, 0.36264321208000183, 0.20785196125507355, 0.07420048117637634, 0.13665802776813507, 0.09254764765501022, 0.1306789219379425, 0.12550167739391327, 0.2261468470096588, 0.23616164922714233, 0.08240598440170288, 0.16113494336605072]
      >
    },
    "mobilenetv20_features_linearbottleneck13_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [-0.03739020228385925, 0.2608458697795868, 0.24877125024795532, 0.1360553354024887, 0.016229359433054924, -0.01019324455410242, 0.031282249838113785, 0.06429833173751831, 0.4442075192928314, 0.12093175202608109, 0.3710899353027344, 0.1019391119480133, 0.05545637756586075, 0.20162232220172882, 0.22989270091056824, -0.056920066475868225, 0.21525850892066956, 0.16917812824249268, 0.29267072677612305, -0.014960628002882004, -0.010674528777599335, 0.034980207681655884, 0.01854848489165306, 0.3039402961730957, 0.3781437873840332, -0.01081379596143961, 0.20649345219135284, 0.17013134062290192, -0.07281264662742615, 0.21421056985855103, 0.1230837032198906, -0.028919575735926628, 0.4143572747707367, 0.4437909722328186, 0.24081338942050934, -0.02267167717218399, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.2845560610294342, 0.08078448474407196, 0.09142519533634186, 0.2485066056251526, 0.12680527567863464, 0.05622509494423866, 0.16645747423171997, 0.09997828304767609, 0.13056667149066925, -0.14148235321044922, 0.20705546438694, 0.15544001758098602, 0.17092552781105042, 0.17278708517551422, 0.11471603810787201, 0.293874055147171, 0.10108096897602081, 0.07669554650783539, 0.0892992690205574, 0.14439816772937775, 0.08082013577222824, 0.09326673299074173, 0.2093047946691513, 0.1999874711036682, 0.09333246201276779, 0.09419676661491394, 0.09989620745182037, 0.10981622338294983, 0.29239147901535034, 0.16639871895313263, 0.15610355138778687, 0.16577696800231934, 0.149077370762825, 0.13844992220401764, 0.16439661383628845, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 576]
        EXLA.Backend
        [-0.003538824152201414, -0.03366425633430481, 0.012803352437913418, 0.09639126062393188, 0.005000478122383356, 3.1017653340692506e-25, 0.06041792035102844, 0.052521444857120514, -0.11556841433048248, 0.09165628999471664, -0.788007378578186, -0.23247763514518738, 0.11072614043951035, -0.4573657214641571, 0.017180858179926872, -0.0022661907132714987, -0.03835657611489296, -0.08881974220275879, 0.04611002653837204, -0.0033401758410036564, -1.4489420573227108e-4, 0.023245586082339287, 0.0013762684538960457, -0.3065164089202881, -0.05932840332388878, -3.603004188335035e-6, -0.10003875195980072, -0.09117628633975983, -0.005860631354153156, -0.09917475283145905, -0.12712354958057404, -0.001667663804255426, -0.7422534227371216, -0.09449290484189987, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 576]
        EXLA.Backend
        [7.354530971497297e-4, 0.003912550862878561, 0.0023845545947551727, 0.04220191761851311, 0.0015303390100598335, 5.115088894877726e-28, 0.018445873633027077, 0.004353330470621586, 0.011003437452018261, 0.0026525913272053003, 0.26168882846832275, 0.06833866983652115, 0.004494751803576946, 0.04838567227125168, 0.0011437759967520833, 2.0979894907213748e-4, 0.002023890847340226, 0.008455991744995117, 0.004080511629581451, 5.398177891038358e-4, 5.815971235278994e-6, 0.0031868095975369215, 0.014848082326352596, 0.007941470481455326, 0.014075135812163353, 5.0474916690745886e-8, 0.012126343324780464, 0.004136150237172842, 5.850492743775249e-4, 0.010100907646119595, 0.007860152050852776, 8.510833868058398e-5, 0.06432832032442093, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck7_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-0.05287745222449303, 0.03939635679125786, 0.008943471126258373, 0.016773579642176628, 0.024517597630620003, -0.2608976662158966, -0.013363493606448174, -0.04500827193260193, -0.0814104974269867, -0.14708103239536285, -0.1051139384508133, -0.03847385570406914, -0.029648277908563614, 0.07605069130659103, -0.042893584817647934, -0.03939994052052498, -0.29128172993659973, -0.03861434757709503, -0.02712533250451088, -0.019969439134001732, -0.07780560106039047, -0.02379690855741501, -0.1076209619641304, -0.028538363054394722, -0.1733534038066864, 0.02688787877559662, -0.05968484282493591, -0.01906135492026806, -0.07085005193948746, -0.04659796133637428, -0.10276751965284348, -0.10676652938127518, -0.05433874577283859, -0.08442538976669312, 0.3125787079334259, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.19615307450294495, 0.13756610453128815, 0.1616729497909546, 0.16233323514461517, 0.1569656878709793, 0.30160748958587646, 0.1933916211128235, 0.18154680728912354, 0.24717095494270325, 0.29326149821281433, 0.2486879527568817, 0.12994907796382904, 0.24061599373817444, 0.13928817212581635, 0.1792965978384018, 0.23451970517635345, 0.23798415064811707, 0.21801051497459412, 0.15246279537677765, 0.20245827734470367, 0.21029138565063477, 0.20429369807243347, 0.26336348056793213, 0.2377043068408966, 0.1920943260192871, 0.14207608997821808, 0.2790174186229706, 0.22336551547050476, 0.2147824764251709, 0.18369099497795105, 0.27606746554374695, 0.28400060534477234, 0.16117443144321442, 0.18702644109725952, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-0.009776907041668892, 0.005012216046452522, -0.047731850296258926, 0.005270622204989195, 0.00557712372392416, -0.00950763002038002, -0.07496191561222076, 0.018082473427057266, 0.0058733015321195126, 0.019821494817733765, 0.008286324329674244, 0.006341803818941116, 0.007210697513073683, -0.021940631791949272, 0.005796361714601517, -0.008641923777759075, 0.012914237566292286, -5.078607355244458e-4, 0.021764321252703667, -0.0017188549973070621, 0.024174479767680168, -0.04216955602169037, -0.041497547179460526, -9.607977117411792e-4, 0.006745839957147837, 0.009343116544187069, 5.745989037677646e-4, 0.005417129956185818, 0.0023382019717246294, 0.002557831583544612, -0.03232024237513542, -0.06076114997267723, 0.0038470544386655092, ...]
      >,
      "var" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [9.27670334931463e-4, 2.7918731211684644e-4, 4.3098966125398874e-4, 0.0011215894483029842, 0.0012489479267969728, 8.453890914097428e-4, 9.005925385281444e-4, 4.5995140681043267e-4, 9.497502469457686e-4, 7.083975942805409e-4, 6.839227862656116e-4, 9.455660620005801e-5, 8.202312747016549e-4, 4.4565973803400993e-4, 3.3621577313169837e-4, 0.0014562587020918727, 6.57685799524188e-4, 7.828701636753976e-4, 2.743843651842326e-4, 0.001106554176658392, 4.3269505840726197e-4, 5.598529824055731e-4, 6.088317022658885e-4, 0.0010905259987339377, 2.4314268375746906e-4, 3.704722039401531e-4, 0.0010972988093271852, 7.521603838540614e-4, 5.540369893424213e-4, 3.005551116075367e-4, 7.429083343595266e-4, 7.382072508335114e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck8_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.03583831712603569, 0.08322400599718094, 0.08347140997648239, 0.16407261788845062, -0.036046866327524185, -0.07233322411775589, -0.007247522007673979, 0.22039948403835297, 0.16788388788700104, 0.011028931476175785, -0.052548106759786606, 0.09673431515693665, -0.009311914443969727, 0.04754066839814186, -0.13346582651138306, 0.0462525449693203, 0.14666330814361572, 0.13988535106182098, 0.08048173785209656, 0.12783612310886383, 0.10406619310379028, 0.061243198812007904, 0.12299603223800659, 0.14473021030426025, 0.15518461167812347, -0.014591885730624199, 0.056073401123285294, 0.1325678676366806, -0.06322228908538818, -0.07645687460899353, -0.04285341501235962, -0.10308938473463058, 0.03268742933869362, -0.04713943228125572, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.12965284287929535, 0.08416002988815308, 0.1387501209974289, 0.0941140204668045, 0.16713306307792664, 0.15179668366909027, 0.1911296546459198, 0.08323527127504349, 0.09111560136079788, 0.142413929104805, 0.12192691117525101, 0.10468165576457977, 0.16143439710140228, 0.14601951837539673, 0.16359323263168335, 0.13607709109783173, 0.0671626552939415, 0.09158993512392044, 0.11166514456272125, 0.12548623979091644, 0.08543179929256439, 0.11259179562330246, 0.06824684143066406, 0.04742749035358429, 0.07571908831596375, 0.129663348197937, 0.09579146653413773, 0.05607251077890396, 0.1601036787033081, 0.15332894027233124, 0.12319595366716385, 0.12570618093013763, 0.11699129641056061, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-2.3525810684077442e-4, 4.89452622787212e-6, -1.9739160779863596e-4, -3.639425849542022e-4, 1.6747551853768528e-4, 4.836112493649125e-4, -8.706151857040823e-5, -3.2933344482444227e-4, -1.4665452181361616e-4, 2.6016379706561565e-4, -7.966592238517478e-5, 3.506306529743597e-5, 1.9936829630751163e-4, 1.982947433134541e-4, -1.9341196457389742e-4, 6.499576265923679e-4, 9.932630928233266e-5, 2.0891131134703755e-4, 1.7995496455114335e-4, 1.8568814266473055e-4, 7.022052159300074e-5, -1.6939605120569468e-4, 1.3075539754936472e-5, 3.1177018536254764e-4, 2.119196142302826e-5, 2.907874295488e-4, 2.8831884264945984e-4, 2.1954343537800014e-4, 8.599817374488339e-5, 2.881138352677226e-4, 5.481694242917001e-4, 3.153556608594954e-4, ...]
      >,
      "var" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.04813230410218239, 0.029002580791711807, 0.04511406272649765, 0.05107865855097771, 0.0974809005856514, 0.0741136446595192, 0.17005448043346405, 0.0427735410630703, 0.04020823910832405, 0.06400398910045624, 0.054310500621795654, 0.0434776172041893, 0.08294457197189331, 0.0641767680644989, 0.09278614073991776, 0.11118675768375397, 0.038928885012865067, 0.03660726547241211, 0.062162626534700394, 0.08595701307058334, 0.03366314619779587, 0.04004068672657013, 0.04966263100504875, 0.052451543509960175, 0.04023684188723564, 0.03982854261994362, 0.04733984172344208, 0.049878451973199844, 0.06715817749500275, 0.08148116618394852, 0.04592667892575264, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck13_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [-0.1382187306880951, 0.030177215114235878, 0.0760938972234726, -0.0933450311422348, -0.062064383178949356, -0.0342678464949131, -0.0573372058570385, 0.0010662958957254887, 0.08649088442325592, 0.12198464572429657, 0.17406152188777924, 0.1671399176120758, 0.10361918807029724, 0.21538838744163513, 0.07684383541345596, -0.1418915092945099, 0.10764788836240768, 0.09419229626655579, 0.14070218801498413, -0.13438130915164948, -0.04806574434041977, -0.0844091922044754, 0.009290370158851147, 0.3029002547264099, 0.049423202872276306, -0.08733383566141129, 0.08972810208797455, 0.10054787248373032, -0.11646779626607895, 0.22748790681362152, 0.14676938951015472, -0.08494674414396286, 0.2620678246021271, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1217566505074501, 0.09843601286411285, 0.07865166664123535, 0.23197029531002045, 0.11923961341381073, 0.005590819288045168, 0.1336921751499176, 0.11597490310668945, 0.11429942399263382, 0.07977456599473953, 0.14655692875385284, 0.10300665348768234, 0.049060989171266556, 0.09101467579603195, 0.08092355728149414, 0.10389447212219238, 0.13443508744239807, 0.11456417292356491, 0.04870767146348953, 0.09771085530519485, 0.026292135939002037, 0.1377914845943451, 0.11884427070617676, 0.0724867507815361, 0.12049250304698944, 0.014269363135099411, 0.12408799678087234, 0.06580907106399536, 0.1170227900147438, 0.056365564465522766, 0.10497873276472092, 0.053536828607320786, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 576]
        EXLA.Backend
        [7.107030910447065e-7, 1.2740528632093628e-7, -4.475431012451736e-7, 3.9450111444239155e-7, 1.2696912108367542e-6, 1.1421681023193742e-7, -1.5435513489592267e-7, 1.8529564727032266e-7, 6.810130912526802e-7, -8.743234616304107e-7, 1.947675400515436e-6, 1.2546711332106497e-6, 7.318482175833196e-7, 1.3930009572504787e-6, -1.3000301350984955e-6, 2.0170767811578116e-7, -9.662721112135841e-8, -9.158228664318813e-8, 1.4846690987724287e-7, -1.0501650660188488e-9, 2.3822990158350876e-7, 3.257930814015708e-7, -2.47027259092647e-7, -1.7559150364832021e-6, 7.632060032847221e-7, 2.672674668247055e-7, 1.0943275583485956e-6, 8.79671517850511e-7, 6.326994821392873e-7, 5.133779268362559e-7, -8.117084462355706e-7, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 576]
        EXLA.Backend
        [0.06938643753528595, 0.07916579395532608, 0.06157204136252403, 0.44644618034362793, 0.11789164692163467, 0.001848552143201232, 0.14640633761882782, 0.10692061483860016, 0.13644804060459137, 0.08004897832870483, 0.48477858304977417, 0.10640207678079605, 0.08469272404909134, 0.15053872764110565, 0.06757758557796478, 0.054888878017663956, 0.14395546913146973, 0.14867006242275238, 0.07945989817380905, 0.041314609348773956, 0.00284106214530766, 0.08277732878923416, 0.1577862948179245, 0.22402343153953552, 0.11104775965213776, 0.008510706946253777, 0.10961493849754333, 0.07491594552993774, 0.06091251224279404, 0.10136393457651138, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck1_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.3692082166671753, 0.019333837553858757, 0.74114990234375, 0.2654167115688324, 0.24005535244941711, 0.04846649616956711, 0.5354508757591248, 0.03959512710571289, 0.2306070178747177, 0.19200822710990906, -0.48113152384757996, 0.006648355163633823, -0.023521361872553825, 0.047236256301403046, -0.015566157177090645, 0.11580895632505417, -0.01523163728415966, -0.0030626191291958094, 0.3765759766101837, -0.01797349564731121, 0.15314938127994537, 0.18993273377418518, 0.3641800582408905, 0.4655008018016815, 0.22646495699882507, 0.0988355278968811, 0.24557138979434967, 0.3279559314250946, 0.0678282305598259, 0.271154522895813, 0.27961465716362, 0.01484939269721508, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.1144828125834465, 0.29103899002075195, 0.19599959254264832, 0.11088918149471283, 0.11658263951539993, 0.2404186874628067, 0.05210668221116066, 0.279130220413208, 0.16053666174411774, 0.10330193489789963, 0.17037393152713776, 0.42853739857673645, 0.2756047546863556, 0.3480927646160126, 0.41298505663871765, 0.3465864956378937, 0.2918163239955902, 0.2078855335712433, 0.19570542871952057, 0.3232663571834564, 0.20460295677185059, 0.21620818972587585, 0.06636220216751099, 0.16237126290798187, 0.11230707168579102, 0.20659472048282623, 0.04134644940495491, 0.16707743704319, 0.18970169126987457, 0.08170202374458313, -0.14146141707897186, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 96]
        EXLA.Backend
        [7.718168199062347e-4, -3.054402768611908e-4, -9.611909626983106e-5, -7.030020788079128e-5, 2.542392030591145e-5, -3.166895767208189e-4, 1.114359765779227e-4, 1.5082783647812903e-4, 4.501911753322929e-4, -4.4038353371433914e-4, -1.635406370041892e-4, 1.342497707810253e-4, 7.245792949106544e-5, 5.361206131055951e-4, 0.0017665454652160406, -4.011236014775932e-4, -1.301867905567633e-5, 7.278626435436308e-4, 4.069282440468669e-4, -2.7848513127537444e-5, 1.5316499047912657e-4, 1.1272649135207757e-4, -1.611710904398933e-4, 4.5594776747748256e-4, 1.0409200331196189e-4, -6.687425047857687e-5, 1.571962347952649e-4, 3.37887613568455e-5, -4.424355283845216e-4, -7.257363176904619e-5, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 96]
        EXLA.Backend
        [0.10025256127119064, 0.17444787919521332, 0.3846805989742279, 0.17531226575374603, 0.19355033338069916, 0.10090412944555283, 0.18123510479927063, 0.20078279078006744, 0.0950746163725853, 0.06799224019050598, 0.14198918640613556, 0.43498390913009644, 0.17685659229755402, 0.4196777641773224, 0.31417417526245117, 0.38556617498397827, 0.09144128859043121, 0.10989599674940109, 0.33278748393058777, 0.11799044907093048, 0.1933024525642395, 0.19662970304489136, 0.13493505120277405, 0.26717954874038696, 0.16789010167121887, 0.15902097523212433, 0.09305445104837418, 0.2792528569698334, 0.12162318080663681, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1379205733537674, 0.1852763593196869, 0.07007614523172379, 0.14786551892757416, 0.08499550819396973, 0.022980524227023125, 0.10445888340473175, -0.03380206599831581, 0.05621454492211342, 0.13879631459712982, 0.1278844028711319, 0.11503273248672485, 0.008648642338812351, 0.07986151427030563, 0.05827014893293381, 0.02635921724140644, 0.16676078736782074, 0.1696624606847763, -0.0049711973406374454, 0.11680029332637787, 0.09153503179550171, -0.09545790404081345, 0.10052905976772308, 0.2198638916015625, 0.09066508710384369, 0.021357741206884384, 0.1243291050195694, 0.13527752459049225, 0.09010029584169388, 0.07285095006227493, 0.09204164147377014, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1116461232304573, 0.08017633110284805, 0.11955763399600983, 0.09390335530042648, 0.09855279326438904, 0.14738331735134125, 0.0888434574007988, 0.13319958746433258, 0.10988157987594604, 0.09758485853672028, 0.10826665163040161, 0.1016610860824585, 0.09888888150453568, 0.12156049907207489, 0.11694491654634476, 0.119488425552845, 0.07615012675523758, 0.06269187480211258, 0.15567438304424286, 0.1064029410481453, 0.11337470263242722, 0.14443783462047577, 0.0840112492442131, 0.09276950359344482, 0.16790293157100677, 0.11844969540834427, 0.0930234044790268, 0.11310834437608719, 0.1162751242518425, 0.1104091927409172, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [3.63793503765919e-7, 1.209602089602413e-7, 5.4800288751266635e-8, 2.552324076532386e-7, 3.2612440747925575e-8, 1.4689121030642127e-8, 1.1114344999896275e-7, -8.35849434110969e-8, 2.2601201976613083e-7, 1.6091770760340296e-7, 3.1186812066152925e-7, -1.1275391287313141e-8, -2.748528871165945e-8, -1.7028881416081276e-7, -1.3741856719207135e-7, -1.3034306789450056e-7, 4.320734205975896e-7, 5.223997732173302e-7, 1.3161741918565895e-8, 1.749855726984606e-7, 6.074458980265263e-8, 4.904541128780693e-7, 2.8845341670802327e-9, 3.752578692228781e-8, -8.265797646345163e-8, -2.6385482243540537e-8, 6.387549156272598e-9, -3.501408230022207e-8, 3.437627640323626e-7, ...]
      >,
      "var" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.02586364559829235, 0.04419300705194473, 0.02695547789335251, 0.027057670056819916, 0.023040946573019028, 0.03236440196633339, 0.021965283900499344, 0.03861037641763687, 0.02072414942085743, 0.020027386024594307, 0.017105598002672195, 0.027752427384257317, 0.025669051334261894, 0.030634647235274315, 0.02309420332312584, 0.026124022901058197, 0.009496251121163368, 0.02531951665878296, 0.04465075209736824, 0.03168340027332306, 0.03208788484334946, 0.03684369474649429, 0.01707514189183712, 0.033794183284044266, 0.05672197788953781, 0.01841931976377964, 0.031393010169267654, 0.02732437662780285, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck15_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-0.08308334648609161, 0.12670446932315826, -0.0052024004980921745, -0.045525070279836655, 0.07931708544492722, 0.005196369253098965, -0.04450135678052902, -0.061325348913669586, 0.07418622076511383, 0.007728567346930504, -0.021894417703151703, 0.04208722338080406, -0.0587613582611084, -0.04845166578888893, -0.07747825980186462, 0.07127270102500916, 0.01305173709988594, -0.0022182995453476906, -0.0679347962141037, 0.03478364646434784, -0.06757770478725433, -0.08550934493541718, -0.14692631363868713, -0.003191823372617364, -0.013460869900882244, -0.06923282146453857, 0.056381601840257645, 0.34844130277633667, -0.034353602677583694, 0.07313910871744156, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.31569308042526245, 0.1642826497554779, 0.22775427997112274, 0.23469418287277222, 0.16511954367160797, 0.13983076810836792, 0.19703467190265656, 0.15652401745319366, 0.14988121390342712, 0.1384556144475937, 0.17888672649860382, 0.2019050568342209, 0.20270013809204102, 0.2201196551322937, 0.22689132392406464, 0.15754374861717224, 0.14099928736686707, 0.13601884245872498, 0.20138750970363617, 0.17591048777103424, 0.194340318441391, 0.17946387827396393, 0.24784566462039948, 0.14162275195121765, 0.2896343469619751, 0.21611864864826202, 0.1647210717201233, 0.11638777703046799, 0.2246612310409546, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.006607928778976202, -0.005489333998411894, 0.001202044659294188, 0.016741983592510223, -0.021014627069234848, 3.9149439544416964e-4, -0.001380614354275167, 0.006143252365291119, -0.022377200424671173, 3.921858442481607e-4, 0.011808911338448524, -0.038100697100162506, -0.0027107521891593933, -0.002684985054656863, -0.003361391369253397, -0.027487538754940033, 0.0032275188714265823, 7.377289002761245e-4, -0.003101653652265668, -0.02393931709229946, -0.0026154580991715193, 4.0451946551911533e-4, 0.019924011081457138, 0.0014483053237199783, 0.012146278284490108, -6.006719195283949e-4, -0.02536875568330288, -0.005276186391711235, ...]
      >,
      "var" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [6.60044839605689e-4, 2.093793300446123e-4, 6.339335814118385e-4, 8.293294231407344e-4, 4.367101937532425e-4, 3.0174001949490048e-5, 1.974338520085439e-4, 5.38308268005494e-5, 5.038247327320278e-4, 1.9694401999004185e-4, 3.8651813520118594e-4, 5.852532922290266e-4, 2.917999809142202e-4, 3.8506241980940104e-4, 3.85848106816411e-4, 4.808907106053084e-4, 8.100475679384544e-5, 6.1370710682240315e-6, 1.8867437029257417e-4, 3.83470905944705e-4, 5.22851652931422e-4, 4.0124301449395716e-5, 4.517717461567372e-4, 7.429129254887812e-6, 0.0013210318284109235, 2.3677274293731898e-4, 4.158295632805675e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck9_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.029621589928865433]
            ],
            [
              [0.03196175768971443]
            ],
            [
              [-0.04775437340140343]
            ],
            [
              [-0.02047593705356121]
            ],
            [
              [0.022681893780827522]
            ],
            [
              [0.0061656758189201355]
            ],
            [
              [-0.020124075934290886]
            ],
            [
              [-0.0419141948223114]
            ],
            [
              [-0.0958712100982666]
            ],
            [
              [-0.03112773969769478]
            ],
            [
              [0.033005211502313614]
            ],
            [
              [-0.0796334370970726]
            ],
            [
              [-0.08826649934053421]
            ],
            [
              [0.008132039569318295]
            ],
            [
              [0.13005898892879486]
            ],
            [
              [0.01112390961498022]
            ],
            [
              [-0.025056127458810806]
            ],
            [
              [-0.021656425669789314]
            ],
            [
              [-0.009619093500077724]
            ],
            [
              [-0.09488093852996826]
            ],
            [
              [-0.03190577030181885]
            ],
            [
              [0.08531086146831512]
            ],
            [
              [0.047572940587997437]
            ],
            [
              [-0.036207523196935654]
            ],
            [
              [-0.1315651535987854]
            ],
            [
              [-0.03746785223484039]
            ],
            [
              [-0.04043968394398689]
            ],
            [
              [-0.1008371040225029]
            ],
            [
              [-7.86909949965775e-4]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][576][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.04216459393501282]
            ],
            [
              [-0.03359765186905861]
            ],
            [
              [0.021302534267306328]
            ],
            [
              [0.06966963410377502]
            ],
            [
              [-0.023086098954081535]
            ],
            [
              [0.02423691935837269]
            ],
            [
              [-0.0027191585395485163]
            ],
            [
              [-0.12812955677509308]
            ],
            [
              [0.10142464935779572]
            ],
            [
              [0.04367773234844208]
            ],
            [
              [-0.04808792844414711]
            ],
            [
              [-0.059124793857336044]
            ],
            [
              [-0.03544275462627411]
            ],
            [
              [-0.10292188823223114]
            ],
            [
              [-0.15540222823619843]
            ],
            [
              [0.08993023633956909]
            ],
            [
              [-0.0018712839810177684]
            ],
            [
              [0.04499407112598419]
            ],
            [
              [0.03717062994837761]
            ],
            [
              [0.1088862419128418]
            ],
            [
              [0.03789613023400307]
            ],
            [
              [0.011467562057077885]
            ],
            [
              [0.022153226658701897]
            ],
            [
              [0.011087901890277863]
            ],
            [
              [0.10830669850111008]
            ],
            [
              [-0.039926622062921524]
            ],
            [
              [0.03126271069049835]
            ],
            [
              [0.11330234259366989]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.06266669183969498, 0.34174367785453796, 0.014498202130198479, 0.09171395748853683, -0.01248596515506506, 0.09248802065849304, 0.03913332521915436, -0.12479757517576218, 0.04716098681092262, 0.2623796761035919, 0.00875355675816536, 0.3866884410381317, 0.34793952107429504, 0.07759835571050644, 0.28590700030326843, 0.34494516253471375, 0.1602264642715454, 0.3233429491519928, -0.08500878512859344, 0.017791876569390297, 0.012547487393021584, 0.07443935424089432, 0.470037043094635, 0.06998748332262039, -0.1145705133676529, 0.35432684421539307, 0.0036894683726131916, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.14569389820098877, 0.08483350276947021, 0.18709848821163177, 0.08500487357378006, 0.0695008710026741, 0.13266052305698395, 0.23066304624080658, 0.1294335573911667, 0.15951289236545563, 0.1612233966588974, 0.04986649379134178, 0.24220621585845947, 0.13059896230697632, 0.2140040248632431, 0.16401590406894684, 0.04375945031642914, 0.1284344643354416, 0.10862768441438675, 0.22237107157707214, 0.20400500297546387, 0.2341395914554596, 0.36944401264190674, 0.1626269817352295, 0.1425088793039322, 0.37626516819000244, 0.10176140815019608, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.04842672124505043, 0.011466434225440025, -0.28272300958633423, 0.04086904227733612, -5.605193857299268e-45, 0.03728378936648369, -0.07810080051422119, -0.0095414062961936, 0.019914837554097176, -0.26291584968566895, 1.200515914815965e-22, 0.4378272294998169, -0.07176417112350464, -0.20427647233009338, -0.40337827801704407, 0.07497201859951019, 0.06511131674051285, 0.02972135879099369, -0.029825769364833832, 0.03816595301032066, 0.008758057840168476, -0.5172790288925171, 0.18057824671268463, 0.03219117224216461, 0.25589102506637573, ...]
      >,
      "var" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.0016082772053778172, 0.008874273858964443, 0.014802338555455208, 0.010124930180609226, 5.605193857299268e-45, 0.0027794779743999243, 0.019753890112042427, 6.709994049742818e-4, 0.004814377520233393, 0.018505984917283058, 3.875728477835647e-26, 0.024166392162442207, 0.011687243357300758, 0.018400687724351883, 0.00805666297674179, 0.003925488330423832, 0.004267431795597076, 0.005561113357543945, 0.00586837250739336, 0.01074590440839529, 0.023998573422431946, 0.010233975946903229, 0.007847330532968044, 0.0024662334471940994, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[144][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.007945280522108078, -0.5135775208473206, 0.642527163028717],
              [0.09571424126625061, -0.4191504120826721, -0.42046400904655457],
              [0.2738749086856842, -1.341597080230713, 1.7941224575042725]
            ]
          ],
          [
            [
              [0.36698344349861145, 0.8006415963172913, 1.0355921983718872],
              [0.2570298910140991, 0.9058852791786194, 0.1435706615447998],
              [0.08509054780006409, -0.29565608501434326, -0.19910597801208496]
            ]
          ],
          [
            [
              [-0.07826036959886551, 1.5095280408859253, -0.7258926033973694],
              [-0.1067059189081192, -0.07988816499710083, -0.4704415500164032],
              [-2.0861942768096924, 0.03404847905039787, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-0.14200253784656525, -0.08138399571180344, 0.020918697118759155, -0.030422303825616837, -0.054856158792972565, 0.0021326385904103518, -0.18490885198116302, -0.03978404402732849, 0.17561353743076324, 0.007802810054272413, -0.05954292416572571, 0.001685986528173089, -0.021205948665738106, 0.09003060311079025, 0.03726910054683685, 0.14319199323654175, 0.15906931459903717, 0.09075694531202316, 0.14122611284255981, 0.19317775964736938, 0.11409415304660797, 0.2862238883972168, 0.21863330900669098, 0.31881415843963623, 0.15870404243469238, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.08955207467079163, 0.12251311540603638, 0.13474930822849274, 0.12169622629880905, 0.10408146679401398, 0.12673817574977875, 0.09429759532213211, 0.004978157579898834, 0.06450905650854111, 0.12292398512363434, 0.008911419659852982, 0.16245806217193604, 0.12929747998714447, 0.07689117640256882, 0.13534259796142578, 0.0939403846859932, 0.07975732535123825, 0.07448924332857132, 0.02857014909386635, 0.05634821206331253, 0.04961058497428894, 0.08511262387037277, 0.1229783222079277, 0.09482692182064056, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 384]
        EXLA.Backend
        [2.307723043486476e-4, 8.844282128848135e-4, 2.083403232973069e-4, 2.6221980806440115e-4, -1.6112398589029908e-4, 3.322429256513715e-4, 1.4785825624130666e-4, 8.785256068222225e-5, -5.301804631017148e-4, 3.369346959516406e-4, 1.6374184633605182e-4, 8.148188935592771e-4, 4.690660280175507e-4, 2.0577036775648594e-4, 7.598484517075121e-4, -1.2347238953225315e-4, -2.2811154485680163e-4, 1.4911795733496547e-4, -1.732072269078344e-5, 1.1792903387686238e-4, 3.792224742937833e-4, 8.303285540023353e-6, 8.0341991269961e-4, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 384]
        EXLA.Backend
        [0.0380934476852417, 0.19316422939300537, 0.16920267045497894, 0.12184684723615646, 0.07579319179058075, 0.16008980572223663, 0.06657657772302628, 0.0039118314161896706, 0.3378544747829437, 0.149428129196167, 0.011692735366523266, 0.34553080797195435, 0.19671620428562164, 0.12972408533096313, 0.27419647574424744, 0.12860825657844543, 0.137705460190773, 0.09913351386785507, 0.12131999433040619, 0.21013996005058289, 0.12196559458971024, 0.39621591567993164, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-2.24758423428284e-7, 6.764846460782792e-8, 9.706442227752632e-8, -2.305486646037025e-8, -2.543328037063475e-7, -8.81496191595943e-8, 8.809087859162901e-8, 2.8762743298216265e-8, 1.4136237780348893e-7, 2.6667668961977142e-8, -2.255288933383781e-8, -2.7744016861674936e-8, 9.968827896500443e-8, 1.3189909964239632e-7, -9.693696512158567e-9, -1.3415302646535565e-7, 1.6071410868789826e-7, -1.5768686978390178e-7, -1.577958812504221e-7, -9.775813225587626e-8, 1.5789125029641582e-7, 3.4908910606645804e-7, -1.5811224329809193e-7, 8.227100778412932e-8, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.187729611992836, 0.34665659070014954, 0.2443280816078186, 0.3304775655269623, 0.22040368616580963, 0.3168470561504364, 0.28643155097961426, 0.16961632668972015, 0.31918713450431824, 0.2289806753396988, 0.21320271492004395, 0.22892144322395325, 0.2310662865638733, 0.17097973823547363, 0.2751845419406891, 0.22792327404022217, 0.26325085759162903, 0.17778971791267395, 0.25794553756713867, 0.22512315213680267, 0.2510114312171936, 0.22274149954319, 0.19815665483474731, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-0.10289469361305237, 0.13421736657619476, 0.3127436637878418, -0.010098671540617943, 0.11501706391572952, 0.07170626521110535, -0.12338873744010925, -0.09662780910730362, 0.10249986499547958, 0.039654918015003204, 0.23144079744815826, -0.02762533910572529, -0.2323276251554489, -0.14802084863185883, -0.17405559122562408, -0.05928152799606323, 0.30801624059677124, 0.03448152542114258, 0.10294467210769653, -0.14233750104904175, 0.042918454855680466, -0.15366889536380768, ...]
      >,
      "var" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.01982937380671501, 0.07223469018936157, 0.03559407591819763, 0.05254514887928963, 0.03399147093296051, 0.05125325173139572, 0.045940618962049484, 0.015112430788576603, 0.0665295273065567, 0.029054192826151848, 0.024110015481710434, 0.028089039027690887, 0.032291218638420105, 0.017801793292164803, 0.04934801161289215, 0.028767650946974754, 0.036390747874975204, 0.020782455801963806, 0.038260187953710556, 0.03192899003624916, 0.0354829803109169, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [-0.0010560675291344523, -6.787155289202929e-4, -7.046025712043047e-4, -9.398604743182659e-4, -8.966192253865302e-4, -2.11737904464826e-4, -5.149749922566116e-4, -0.0010324727045372128, -7.925425306893885e-4, -9.090156527236104e-5, -4.453391011338681e-4, -9.805577574297786e-4, -7.45584664400667e-4, -7.640189724043012e-4, -0.0011470100143924356, -3.8889769348315895e-4]
      >,
      "gamma" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.5097491145133972, 0.6192775368690491, 0.657399594783783, 0.5558041930198669, 0.6164559721946716, 0.6381960511207581, 0.6152222156524658, 0.4946329891681671, 0.5616724491119385, 0.6567564606666565, 0.5526584982872009, 0.42007970809936523, 0.4637764096260071, 0.5626768469810486, 0.6189385056495667, 0.6354565620422363]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 16]
        EXLA.Backend
        [0.02104923129081726, -0.47712182998657227, -0.46247294545173645, 1.7340073585510254, 0.3934646546840668, -0.8456623554229736, 1.2444655895233154, -0.48179635405540466, 0.8393679857254028, -2.194596767425537, 2.2858681678771973, 0.04296046495437622, -1.1514852046966553, -0.37791600823402405, -3.0484585762023926, -0.7847945094108582]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 16]
        EXLA.Backend
        [0.015505708754062653, 0.05742154270410538, 0.05579017847776413, 0.01873898319900036, 0.03846978396177292, 0.09190379083156586, 0.05365859344601631, 0.04166821017861366, 0.03757927194237709, 0.14188292622566223, 0.04087332263588905, 0.06089872494339943, 0.03475071117281914, 0.05605774745345116, 0.0412914864718914, 0.05219418555498123]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.09069628268480301, 0.011556080542504787, -0.044405367225408554],
              [0.1974998563528061, -0.05144780874252319, -0.15406860411167145],
              [0.03294772282242775, -0.08055609464645386, -0.10401368886232376]
            ]
          ],
          [
            [
              [0.004530614707618952, -0.016135336831212044, 0.010837683454155922],
              [-0.013911036774516106, -0.17941752076148987, 0.0742872878909111],
              [0.03701423481106758, 0.17489252984523773, 0.049539677798748016]
            ]
          ],
          [
            [
              [-0.03442509472370148, -0.04338786378502846, -8.070958429016173e-4],
              [-0.06931383162736893, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck1_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][16][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.26627230644226074]
            ],
            [
              [0.1223592683672905]
            ],
            [
              [-0.17675793170928955]
            ],
            [
              [-0.14175795018672943]
            ],
            [
              [-0.04847275838255882]
            ],
            [
              [0.02572016417980194]
            ],
            [
              [0.18247880041599274]
            ],
            [
              [-0.22272470593452454]
            ],
            [
              [-0.0375950001180172]
            ],
            [
              [-0.12036115676164627]
            ],
            [
              [-0.42806223034858704]
            ],
            [
              [0.2049194574356079]
            ],
            [
              [0.025416148826479912]
            ],
            [
              [-0.20868347585201263]
            ],
            [
              [-0.00869692862033844]
            ],
            [
              [0.10014521330595016]
            ]
          ],
          [
            [
              [0.1901395469903946]
            ],
            [
              [-0.15854698419570923]
            ],
            [
              [-0.05858785659074783]
            ],
            [
              [-0.016011832281947136]
            ],
            [
              [-0.07135247439146042]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-6.408389890566468e-5, -1.0128262510988861e-4, -4.3479518353706226e-5, -1.589354214956984e-4, -7.557999197160825e-5, -2.3196584152174182e-5, -2.216532493548584e-6, 9.525781933916733e-6, -8.600373985245824e-5, -2.3008385323919356e-4, -2.3174902889877558e-4, -2.012878976529464e-4, -9.838923870120198e-5, -1.22729703434743e-4, -9.646145917940885e-5, -2.797471643134486e-5, -1.039204653352499e-4, -1.2249314750079066e-4, -8.16335596027784e-5, -1.541685633128509e-4, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.5065965056419373, 0.44786974787712097, 0.5534567832946777, 0.3781253695487976, 0.44674113392829895, 0.28593751788139343, 0.5029821395874023, 0.19650447368621826, 0.3918337821960449, 0.33050450682640076, 0.31709131598472595, 0.1402808576822281, 0.4198921024799347, 0.02611447498202324, 0.5337468981742859, 0.49816271662712097, 0.32386210560798645, 0.40961867570877075, 0.3363959491252899, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-0.49676281213760376, -0.5412184000015259, 0.2777172029018402, 0.30402812361717224, -0.038367077708244324, 0.19269412755966187, -0.20861922204494476, 0.3330950438976288, -0.35258007049560547, -0.0413212850689888, -0.44278526306152344, 0.43556883931159973, -0.007557356264442205, -0.049789782613515854, -0.1810888946056366, -0.1666330248117447, 0.15477608144283295, 0.1351827085018158, ...]
      >,
      "var" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.0421639047563076, 0.03801903501152992, 0.04582618921995163, 0.028266852721571922, 0.025787167251110077, 0.026369646191596985, 0.03499341756105423, 0.02089923992753029, 0.02381058968603611, 0.04512317106127739, 0.018083490431308746, 0.02251904457807541, 0.023619014769792557, 0.005619029048830271, 0.04942920058965683, 0.02654757723212242, 0.02269740216434002, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.0020327651873230934, -0.03829871118068695, -0.013563335873186588],
              [-0.07023981958627701, -0.0530773401260376, -0.05999664217233658],
              [-0.02172546088695526, -0.05160605534911156, -0.02340804785490036]
            ]
          ],
          [
            [
              [-0.06414368003606796, -0.09891236573457718, 0.09147416055202484],
              [-0.035863131284713745, -0.026917265728116035, 0.07736452668905258],
              [0.08244364708662033, 0.10789858549833298, -0.16784991323947906]
            ]
          ],
          [
            [
              [0.13550637662410736, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.0561985969543457, -0.08534234762191772, -0.0728735700249672],
              [0.20427849888801575, -0.03384411707520485, -0.18976956605911255],
              [0.10962294042110443, 0.006988390814512968, -0.03254517540335655]
            ]
          ],
          [
            [
              [-0.007281317375600338, -0.10424316674470901, -0.0011419359361752868],
              [-0.0659799799323082, -0.2147035449743271, -0.10787361860275269],
              [0.004407626111060381, 0.44921964406967163, 0.06901120394468307]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.0373462438583374]
            ],
            [
              [-0.07842571288347244]
            ],
            [
              [-0.0013282403815537691]
            ],
            [
              [-0.01943795755505562]
            ],
            [
              [-0.012809362262487411]
            ],
            [
              [-0.03409231826663017]
            ],
            [
              [0.00997056532651186]
            ],
            [
              [-0.015080451965332031]
            ],
            [
              [0.04131586477160454]
            ],
            [
              [-0.06238787993788719]
            ],
            [
              [0.0022582856472581625]
            ],
            [
              [0.025756919756531715]
            ],
            [
              [-0.07886741310358047]
            ],
            [
              [-0.027919018641114235]
            ],
            [
              [-0.01644008234143257]
            ],
            [
              [-0.03873310983181]
            ],
            [
              [0.05929648503661156]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck15_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.008749016560614109, 0.020535919815301895, 6.699140067212284e-4],
              [0.00949698593467474, 0.23564079403877258, 0.038606781512498856],
              [-0.057412438094615936, 0.02893063984811306, -0.07270332425832748]
            ]
          ],
          [
            [
              [-0.011751603335142136, -0.16229715943336487, -0.03423398733139038],
              [0.0071285502053797245, 0.25753602385520935, 0.014275305904448032],
              [-0.04593678563833237, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.03408721089363098]
            ],
            [
              [-0.017362290993332863]
            ],
            [
              [0.01995295286178589]
            ],
            [
              [-0.051705457270145416]
            ],
            [
              [0.0689198449254036]
            ],
            [
              [0.08685684949159622]
            ],
            [
              [-0.09933704882860184]
            ],
            [
              [0.07178854197263718]
            ],
            [
              [0.044526729732751846]
            ],
            [
              [-0.018590005114674568]
            ],
            [
              [0.013196618296205997]
            ],
            [
              [-0.0795169249176979]
            ],
            [
              [-0.10800802707672119]
            ],
            [
              [-0.042785294353961945]
            ],
            [
              [0.028526680544018745]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.031621839851140976, 0.18290314078330994, 0.040609210729599],
              [-0.19228996336460114, 0.031516727060079575, 0.09460115432739258],
              [-0.026457147672772408, 0.05563095211982727, 0.0026122310664504766]
            ]
          ],
          [
            [
              [-0.00655049504712224, -0.004564541857689619, -0.01380461547523737],
              [-0.002132135210558772, -0.17951297760009766, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.048515044152736664]
            ],
            [
              [-0.03756922110915184]
            ],
            [
              [0.0977545902132988]
            ],
            [
              [0.03282283991575241]
            ],
            [
              [-0.05638144165277481]
            ],
            [
              [0.03332485631108284]
            ],
            [
              [0.035065557807683945]
            ],
            [
              [0.047478508204221725]
            ],
            [
              [0.034366436302661896]
            ],
            [
              [0.005331412889063358]
            ],
            [
              [0.015283018350601196]
            ],
            [
              [-0.045142702758312225]
            ],
            [
              [-0.018845317885279655]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [2.1335902999908285e-7, 3.535939754328865e-7, 5.073925990473072e-7, 5.08670098042785e-7, 3.3390904263796983e-7, 4.946827516505437e-7, 4.795706445293035e-7, 3.743390948329761e-7, 5.272184466775798e-7, 5.93354513966915e-7, 2.1419212714590685e-7, 6.500247309304541e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.2979993522167206, 0.3140060007572174, 0.2445983588695526, 0.2865447998046875, 0.35691216588020325, 0.3592449426651001, 0.27936986088752747, 0.25922414660453796, 0.2633695900440216, 0.2901868522167206, 0.2989289462566376, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 96]
        EXLA.Backend
        [-0.0291385967284441, -0.5151070356369019, 0.05079524964094162, -0.3128364384174347, -0.02873671054840088, 0.21749934554100037, 0.10778770595788956, -0.23229928314685822, -0.46638214588165283, 0.027039799839258194, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 96]
        EXLA.Backend
        [0.10976872593164444, 0.09856210649013519, 0.06113961338996887, 0.09309827536344528, 0.19614680111408234, 0.13435077667236328, 0.0629381388425827, 0.07282107323408127, 0.058386534452438354, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.03325330838561058, 0.027691686525940895, -0.044400639832019806],
              [0.21181674301624298, -0.09588920325040817, -0.12353076785802841],
              [0.050757963210344315, -0.0020931188482791185, -0.04063280299305916]
            ]
          ],
          [
            [
              [0.025990355759859085, -0.034638673067092896, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck12_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][96][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.020311787724494934]
            ],
            [
              [0.00342985475435853]
            ],
            [
              [-0.04346230626106262]
            ],
            [
              [-0.016626032069325447]
            ],
            [
              [-0.012114123441278934]
            ],
            [
              [-0.012785198166966438]
            ],
            [
              [0.056096695363521576]
            ],
            [
              [0.09021388739347458]
            ],
            [
              [0.16139942407608032]
            ],
            [
              [0.06402886658906937]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [1.037217140197754, 1.2865849733352661, 0.7645959258079529, 0.6374938488006592, 1.069508671760559, 0.6837575435638428, 0.9550023674964905, 0.9839419722557068, 1.1442939043045044, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.2217620462179184, 0.11614178121089935, 0.1586529165506363, 0.15199987590312958, 0.23966756463050842, 0.17114117741584778, 0.18509696424007416, 0.18513236939907074, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 32]
        EXLA.Backend
        [0.6255861520767212, -0.0877540335059166, 1.0720124244689941, 2.3868625164031982, -5.5901641845703125, 3.4052634239196777, 7.124544143676758, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 32]
        EXLA.Backend
        [0.31226837635040283, 0.008983085863292217, 0.010836457833647728, 0.33667588233947754, 2.1207056045532227, 0.36680367588996887, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck6_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.014363789930939674, 0.13237811625003815, 0.09334706515073776],
              [0.12919868528842926, -0.35225632786750793, -0.1466313749551773],
              [0.06540388613939285, -0.08005677908658981, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [-2.8074174451830913e-7, -1.2177669361790322e-7, -1.7438522093016218e-7, 2.43050635617692e-7, 1.9082594349129067e-7, -2.2253024667406862e-7, -3.5373562923268764e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [0.25085631012916565, 0.247349813580513, 0.2442786693572998, 0.2477291077375412, 0.25060394406318665, 0.2502780854701996, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[channels: 320]
        EXLA.Backend
        [-0.11591224372386932, -0.06157569959759712, -0.160318985581398, 0.18986251950263977, -0.1408201903104782, ...]
      >,
      "var" => #Nx.Tensor<
        f32[channels: 320]
        EXLA.Backend
        [0.0080089271068573, 0.01121740322560072, 0.008632301352918148, 0.0062346188351511955, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck10_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][384][1][1]
        EXLA.Backend
        [
          [
            [
              [0.01824713684618473]
            ],
            [
              [0.08710300177335739]
            ],
            [
              [-0.10191874951124191]
            ],
            [
              [-0.07710960507392883]
            ],
            [
              [-0.042343202978372574]
            ],
            [
              [-0.014866536483168602]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][64][1][1]
        EXLA.Backend
        [
          [
            [
              [0.015875784680247307]
            ],
            [
              [0.00680553587153554]
            ],
            [
              [0.09295158088207245]
            ],
            [
              [-0.05411294475197792]
            ],
            [
              [0.026204092428088188]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][3][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.11137320101261139, -0.22142910957336426, 0.10717468708753586],
              [4.23799006966874e-4, ...],
              ...
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [3.9526614159512974e-7, 5.761955890193349e-7, 7.01267651948001e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.21182486414909363, 0.20167435705661774, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [-0.039972804486751556, ...]
      >,
      "var" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [...]
      >
    },
    "mobilenetv20_features_linearbottleneck6_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][192][1][1]
        EXLA.Backend
        [
          [
            [
              [0.07365027070045471]
            ],
            [
              [-0.03652043268084526]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][144][1][1]
        EXLA.Backend
        [
          [
            [
              [0.02765008620917797]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][32][1][1]
        EXLA.Backend
        [
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_conv0_fwd" => %{...},
    ...
  }
}
eval_model = model |> Axon.sigmoid()

eval_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)
Batch: 23, accuracy: 0.9709823
%{
  0 => %{
    "accuracy" => #Nx.Tensor<
      f32
      EXLA.Backend
      0.9709823131561279
    >
  }
}