Training a Simple LSTM
In this tutorial we will go over using a recurrent neural network to classify clockwise and anticlockwise spirals. By the end of this tutorial you will be able to:
Create custom Lux models.
Become familiar with the Lux recurrent neural network API.
Training using Optimisers.jl and Zygote.jl.
Package Imports
Note: If you wish to use AutoZygote() for automatic differentiation, add Zygote to your project dependencies and include using Zygote.
julia
using ADTypes, Lux, JLD2, MLUtils, Optimisers, Printf, Reactant, RandomDataset
We will use MLUtils to generate 500 (noisy) clockwise and 500 (noisy) anticlockwise spirals. Using this data we will create a MLUtils.DataLoader. Our dataloader will give us sequences of size 2 × seq_len × batch_size and we need to predict a binary value whether the sequence is clockwise or anticlockwise.
julia
function get_dataloaders(; dataset_size=1000, sequence_length=50)
# Create the spirals
data = [MLUtils.Datasets.make_spiral(sequence_length) for _ in 1:dataset_size]
# Get the labels
labels = vcat(repeat([0.0f0], dataset_size ÷ 2), repeat([1.0f0], dataset_size ÷ 2))
clockwise_spirals = [
reshape(d[1][:, 1:sequence_length], :, sequence_length, 1) for
d in data[1:(dataset_size ÷ 2)]
]
anticlockwise_spirals = [
reshape(d[1][:, (sequence_length + 1):end], :, sequence_length, 1) for
d in data[((dataset_size ÷ 2) + 1):end]
]
x_data = Float32.(cat(clockwise_spirals..., anticlockwise_spirals...; dims=3))
# Split the dataset
(x_train, y_train), (x_val, y_val) = splitobs((x_data, labels); at=0.8, shuffle=true)
# Create DataLoaders
return (
# Use DataLoader to automatically minibatch and shuffle the data
DataLoader(
collect.((x_train, y_train)); batchsize=128, shuffle=true, partial=false
),
# Don't shuffle the validation data
DataLoader(collect.((x_val, y_val)); batchsize=128, shuffle=false, partial=false),
)
endget_dataloaders (generic function with 1 method)Creating a Classifier
We will be extending the Lux.AbstractLuxContainerLayer type for our custom model since it will contain a lstm block and a classifier head.
We pass the fieldnames lstm_cell and classifier to the type to ensure that the parameters and states are automatically populated and we don't have to define Lux.initialparameters and Lux.initialstates.
To understand more about container layers, please look at Container Layer.
julia
struct SpiralClassifier{L,C} <: AbstractLuxContainerLayer{(:lstm_cell, :classifier)}
lstm_cell::L
classifier::C
endWe won't define the model from scratch but rather use the Lux.LSTMCell and Lux.Dense.
julia
function SpiralClassifier(in_dims, hidden_dims, out_dims)
return SpiralClassifier(
LSTMCell(in_dims => hidden_dims), Dense(hidden_dims => out_dims, sigmoid)
)
endMain.var"##230".SpiralClassifierWe can use default Lux blocks – Recurrence(LSTMCell(in_dims => hidden_dims) – instead of defining the following. But let's still do it for the sake of it.
Now we need to define the behavior of the Classifier when it is invoked.
julia
function (s::SpiralClassifier)(
x::AbstractArray{T,3}, ps::NamedTuple, st::NamedTuple
) where {T}
# First we will have to run the sequence through the LSTM Cell
# The first call to LSTM Cell will create the initial hidden state
# See that the parameters and states are automatically populated into a field called
# `lstm_cell` We use `eachslice` to get the elements in the sequence without copying,
# and `Iterators.peel` to split out the first element for LSTM initialization.
x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2)))
(y, carry), st_lstm = s.lstm_cell(x_init, ps.lstm_cell, st.lstm_cell)
# Now that we have the hidden state and memory in `carry` we will pass the input and
# `carry` jointly
for x in x_rest
(y, carry), st_lstm = s.lstm_cell((x, carry), ps.lstm_cell, st_lstm)
end
# After running through the sequence we will pass the output through the classifier
y, st_classifier = s.classifier(y, ps.classifier, st.classifier)
# Finally remember to create the updated state
st = merge(st, (classifier=st_classifier, lstm_cell=st_lstm))
return vec(y), st
endUsing the @compact API
We can also define the model using the Lux.@compact API, which is a more concise way of defining models. This macro automatically handles the boilerplate code for you and as such we recommend this way of defining custom layers
julia
function SpiralClassifierCompact(in_dims, hidden_dims, out_dims)
lstm_cell = LSTMCell(in_dims => hidden_dims)
classifier = Dense(hidden_dims => out_dims, sigmoid)
return @compact(; lstm_cell, classifier) do x::AbstractArray{T,3} where {T}
x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2)))
y, carry = lstm_cell(x_init)
for x in x_rest
y, carry = lstm_cell((x, carry))
end
@return vec(classifier(y))
end
endSpiralClassifierCompact (generic function with 1 method)Defining Accuracy, Loss and Optimiser
Now let's define the binarycrossentropy loss. Typically it is recommended to use logitbinarycrossentropy since it is more numerically stable, but for the sake of simplicity we will use binarycrossentropy.
julia
const lossfn = BinaryCrossEntropyLoss()
function compute_loss(model, ps, st, (x, y))
ŷ, st_ = model(x, ps, st)
loss = lossfn(ŷ, y)
return loss, st_, (; y_pred=ŷ)
end
matches(y_pred, y_true) = sum((y_pred .> 0.5f0) .== y_true)
accuracy(y_pred, y_true) = matches(y_pred, y_true) / length(y_pred)accuracy (generic function with 1 method)Training the Model
julia
function main(model_type)
dev = reactant_device()
cdev = cpu_device()
# Get the dataloaders
train_loader, val_loader = dev(get_dataloaders())
# Create the model
model = model_type(2, 8, 1)
ps, st = dev(Lux.setup(Random.default_rng(), model))
train_state = Training.TrainState(model, ps, st, Adam(0.01f0))
model_compiled = if dev isa ReactantDevice
@compile model(first(train_loader)[1], ps, Lux.testmode(st))
else
model
end
ad = dev isa ReactantDevice ? AutoEnzyme() : AutoZygote()
for epoch in 1:25
# Train the model
total_loss = 0.0f0
total_samples = 0
for (x, y) in train_loader
(_, loss, _, train_state) = Training.single_train_step!(
ad, lossfn, (x, y), train_state
)
total_loss += loss * length(y)
total_samples += length(y)
end
@printf "Epoch [%3d]: Loss %4.5f\n" epoch (total_loss / total_samples)
# Validate the model
total_acc = 0.0f0
total_loss = 0.0f0
total_samples = 0
st_ = Lux.testmode(train_state.states)
for (x, y) in val_loader
ŷ, st_ = model_compiled(x, train_state.parameters, st_)
ŷ, y = cdev(ŷ), cdev(y)
total_acc += accuracy(ŷ, y) * length(y)
total_loss += lossfn(ŷ, y) * length(y)
total_samples += length(y)
end
@printf "Validation:\tLoss %4.5f\tAccuracy %4.5f\n" (total_loss / total_samples) (
total_acc / total_samples
)
end
return cpu_device()((train_state.parameters, train_state.states))
end
ps_trained, st_trained = main(SpiralClassifier)┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
└ @ LuxCore /var/lib/buildkite-agent/builds/gpuci-10/julialang/lux-dot-jl/lib/LuxCore/src/LuxCore.jl:18
2025-05-08 13:32:16.382854: I external/xla/xla/service/service.cc:152] XLA service 0x2369b6f0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2025-05-08 13:32:16.383183: I external/xla/xla/service/service.cc:160] StreamExecutor device (0): NVIDIA A100-PCIE-40GB MIG 1g.5gb, Compute Capability 8.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1746711136.383890 3123193 se_gpu_pjrt_client.cc:1026] Using BFC allocator.
I0000 00:00:1746711136.383965 3123193 gpu_helpers.cc:136] XLA backend allocating 3825205248 bytes on device 0 for BFCAllocator.
I0000 00:00:1746711136.383997 3123193 gpu_helpers.cc:177] XLA backend will use up to 1275068416 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1746711136.399874 3123193 cuda_dnn.cc:529] Loaded cuDNN version 90400
E0000 00:00:1746711197.044877 3123193 buffer_comparator.cc:145] Difference at 16: 0, expected 0.966326
E0000 00:00:1746711197.044930 3123193 buffer_comparator.cc:145] Difference at 17: 0, expected 0.955446
E0000 00:00:1746711197.044941 3123193 buffer_comparator.cc:145] Difference at 18: 0, expected 0.522552
E0000 00:00:1746711197.044944 3123193 buffer_comparator.cc:145] Difference at 19: 0, expected 0.554959
E0000 00:00:1746711197.044947 3123193 buffer_comparator.cc:145] Difference at 20: 0, expected 0.833471
E0000 00:00:1746711197.044950 3123193 buffer_comparator.cc:145] Difference at 21: 0, expected 0.404081
E0000 00:00:1746711197.044953 3123193 buffer_comparator.cc:145] Difference at 22: 0, expected 0.289287
E0000 00:00:1746711197.044956 3123193 buffer_comparator.cc:145] Difference at 23: 0, expected 0.732437
E0000 00:00:1746711197.044958 3123193 buffer_comparator.cc:145] Difference at 24: 0, expected 1.02391
E0000 00:00:1746711197.044961 3123193 buffer_comparator.cc:145] Difference at 25: 0, expected 0.647103
2025-05-08 13:33:17.044972: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.047637 3123193 buffer_comparator.cc:145] Difference at 16: 0, expected 0.966326
E0000 00:00:1746711197.047649 3123193 buffer_comparator.cc:145] Difference at 17: 0, expected 0.955446
E0000 00:00:1746711197.047653 3123193 buffer_comparator.cc:145] Difference at 18: 0, expected 0.522552
E0000 00:00:1746711197.047657 3123193 buffer_comparator.cc:145] Difference at 19: 0, expected 0.554959
E0000 00:00:1746711197.047660 3123193 buffer_comparator.cc:145] Difference at 20: 0, expected 0.833471
E0000 00:00:1746711197.047663 3123193 buffer_comparator.cc:145] Difference at 21: 0, expected 0.404081
E0000 00:00:1746711197.047666 3123193 buffer_comparator.cc:145] Difference at 22: 0, expected 0.289287
E0000 00:00:1746711197.047669 3123193 buffer_comparator.cc:145] Difference at 23: 0, expected 0.732437
E0000 00:00:1746711197.047671 3123193 buffer_comparator.cc:145] Difference at 24: 0, expected 1.02391
E0000 00:00:1746711197.047674 3123193 buffer_comparator.cc:145] Difference at 25: 0, expected 0.647103
2025-05-08 13:33:17.047680: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.050069 3123193 buffer_comparator.cc:145] Difference at 16: 0, expected 0.966326
E0000 00:00:1746711197.050083 3123193 buffer_comparator.cc:145] Difference at 17: 0, expected 0.955446
E0000 00:00:1746711197.050088 3123193 buffer_comparator.cc:145] Difference at 18: 0, expected 0.522552
E0000 00:00:1746711197.050092 3123193 buffer_comparator.cc:145] Difference at 19: 0, expected 0.554959
E0000 00:00:1746711197.050096 3123193 buffer_comparator.cc:145] Difference at 20: 0, expected 0.833471
E0000 00:00:1746711197.050099 3123193 buffer_comparator.cc:145] Difference at 21: 0, expected 0.404081
E0000 00:00:1746711197.050102 3123193 buffer_comparator.cc:145] Difference at 22: 0, expected 0.289287
E0000 00:00:1746711197.050105 3123193 buffer_comparator.cc:145] Difference at 23: 0, expected 0.732437
E0000 00:00:1746711197.050109 3123193 buffer_comparator.cc:145] Difference at 24: 0, expected 1.02391
E0000 00:00:1746711197.050112 3123193 buffer_comparator.cc:145] Difference at 25: 0, expected 0.647103
2025-05-08 13:33:17.050117: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.052505 3123193 buffer_comparator.cc:145] Difference at 16: 0, expected 0.966326
E0000 00:00:1746711197.052518 3123193 buffer_comparator.cc:145] Difference at 17: 0, expected 0.955446
E0000 00:00:1746711197.052522 3123193 buffer_comparator.cc:145] Difference at 18: 0, expected 0.522552
E0000 00:00:1746711197.052525 3123193 buffer_comparator.cc:145] Difference at 19: 0, expected 0.554959
E0000 00:00:1746711197.052528 3123193 buffer_comparator.cc:145] Difference at 20: 0, expected 0.833471
E0000 00:00:1746711197.052531 3123193 buffer_comparator.cc:145] Difference at 21: 0, expected 0.404081
E0000 00:00:1746711197.052534 3123193 buffer_comparator.cc:145] Difference at 22: 0, expected 0.289287
E0000 00:00:1746711197.052537 3123193 buffer_comparator.cc:145] Difference at 23: 0, expected 0.732437
E0000 00:00:1746711197.052540 3123193 buffer_comparator.cc:145] Difference at 24: 0, expected 1.02391
E0000 00:00:1746711197.052543 3123193 buffer_comparator.cc:145] Difference at 25: 0, expected 0.647103
2025-05-08 13:33:17.052548: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.054926 3123193 buffer_comparator.cc:145] Difference at 32: 0, expected 0.904315
E0000 00:00:1746711197.054938 3123193 buffer_comparator.cc:145] Difference at 33: 0, expected 1.02658
E0000 00:00:1746711197.054942 3123193 buffer_comparator.cc:145] Difference at 34: 0, expected 0.512492
E0000 00:00:1746711197.054946 3123193 buffer_comparator.cc:145] Difference at 35: 0, expected 0.434209
E0000 00:00:1746711197.054949 3123193 buffer_comparator.cc:145] Difference at 36: 0, expected 0.218704
E0000 00:00:1746711197.054951 3123193 buffer_comparator.cc:145] Difference at 37: 0, expected 0.551313
E0000 00:00:1746711197.054954 3123193 buffer_comparator.cc:145] Difference at 38: 0, expected 1.10187
E0000 00:00:1746711197.054957 3123193 buffer_comparator.cc:145] Difference at 39: 0, expected 0.347384
E0000 00:00:1746711197.054960 3123193 buffer_comparator.cc:145] Difference at 40: 0, expected 0.789874
E0000 00:00:1746711197.054963 3123193 buffer_comparator.cc:145] Difference at 41: 0, expected 0.204116
2025-05-08 13:33:17.054968: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.057565 3123193 buffer_comparator.cc:145] Difference at 32: 0, expected 0.904315
E0000 00:00:1746711197.057577 3123193 buffer_comparator.cc:145] Difference at 33: 0, expected 1.02658
E0000 00:00:1746711197.057581 3123193 buffer_comparator.cc:145] Difference at 34: 0, expected 0.512492
E0000 00:00:1746711197.057584 3123193 buffer_comparator.cc:145] Difference at 35: 0, expected 0.434209
E0000 00:00:1746711197.057587 3123193 buffer_comparator.cc:145] Difference at 36: 0, expected 0.218704
E0000 00:00:1746711197.057590 3123193 buffer_comparator.cc:145] Difference at 37: 0, expected 0.551313
E0000 00:00:1746711197.057593 3123193 buffer_comparator.cc:145] Difference at 38: 0, expected 1.10187
E0000 00:00:1746711197.057596 3123193 buffer_comparator.cc:145] Difference at 39: 0, expected 0.347384
E0000 00:00:1746711197.057599 3123193 buffer_comparator.cc:145] Difference at 40: 0, expected 0.789874
E0000 00:00:1746711197.057602 3123193 buffer_comparator.cc:145] Difference at 41: 0, expected 0.204116
2025-05-08 13:33:17.057606: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.060000 3123193 buffer_comparator.cc:145] Difference at 32: 0, expected 0.904315
E0000 00:00:1746711197.060013 3123193 buffer_comparator.cc:145] Difference at 33: 0, expected 1.02658
E0000 00:00:1746711197.060017 3123193 buffer_comparator.cc:145] Difference at 34: 0, expected 0.512492
E0000 00:00:1746711197.060020 3123193 buffer_comparator.cc:145] Difference at 35: 0, expected 0.434209
E0000 00:00:1746711197.060023 3123193 buffer_comparator.cc:145] Difference at 36: 0, expected 0.218704
E0000 00:00:1746711197.060026 3123193 buffer_comparator.cc:145] Difference at 37: 0, expected 0.551313
E0000 00:00:1746711197.060029 3123193 buffer_comparator.cc:145] Difference at 38: 0, expected 1.10187
E0000 00:00:1746711197.060032 3123193 buffer_comparator.cc:145] Difference at 39: 0, expected 0.347384
E0000 00:00:1746711197.060035 3123193 buffer_comparator.cc:145] Difference at 40: 0, expected 0.789874
E0000 00:00:1746711197.060037 3123193 buffer_comparator.cc:145] Difference at 41: 0, expected 0.204116
2025-05-08 13:33:17.060042: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.062624 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.062637 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.062641 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.062644 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.062647 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.062650 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.062653 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.062656 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.062658 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.062661 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.062666: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.065058 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.065071 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.065075 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.065078 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.065081 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.065084 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.065087 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.065090 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.065093 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.065096 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.065101: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.067491 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.067504 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.067508 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.067512 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.067515 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.067518 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.067521 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.067524 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.067527 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.067530 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.067534: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.069926 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.069939 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.069943 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.069946 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.069949 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.069952 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.069955 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.069958 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.069961 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.069964 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.069968: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.072355 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.072367 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.072371 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.072374 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.072377 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.072380 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.072383 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.072386 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.072389 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.072392 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.072397: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.074787 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.074799 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.074803 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.074806 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.074809 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.074812 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.074815 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.074819 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.074822 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.074825 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.074830: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.077229 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.077242 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.077246 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.077249 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.077252 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.077255 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.077257 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.077260 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.077263 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.077266 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.077271: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.079668 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.079681 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.079685 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.079688 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.079690 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.079693 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.079696 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.079699 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.079702 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.079705 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.079710: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.082197 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.082209 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.082213 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.082216 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.082219 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.082222 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.082225 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.082228 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.082231 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.082234 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.082239: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.084635 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.084648 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.084652 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.084655 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.084657 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.084660 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.084663 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.084666 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.084669 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.084672 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.084677: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.087239 3123193 buffer_comparator.cc:145] Difference at 64: 0, expected 0.629991
E0000 00:00:1746711197.087251 3123193 buffer_comparator.cc:145] Difference at 65: 0, expected 0.54577
E0000 00:00:1746711197.087255 3123193 buffer_comparator.cc:145] Difference at 66: 0, expected 0.316298
E0000 00:00:1746711197.087258 3123193 buffer_comparator.cc:145] Difference at 67: 0, expected 0.438545
E0000 00:00:1746711197.087261 3123193 buffer_comparator.cc:145] Difference at 68: 0, expected 0.523314
E0000 00:00:1746711197.087264 3123193 buffer_comparator.cc:145] Difference at 69: 0, expected 0.83106
E0000 00:00:1746711197.087266 3123193 buffer_comparator.cc:145] Difference at 70: 0, expected 0.617399
E0000 00:00:1746711197.087269 3123193 buffer_comparator.cc:145] Difference at 71: 0, expected 0.692252
E0000 00:00:1746711197.087272 3123193 buffer_comparator.cc:145] Difference at 72: 0, expected 0.185378
E0000 00:00:1746711197.087275 3123193 buffer_comparator.cc:145] Difference at 73: 0, expected 0.689502
2025-05-08 13:33:17.087280: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.089672 3123193 buffer_comparator.cc:145] Difference at 128: 0, expected 1.00573
E0000 00:00:1746711197.089684 3123193 buffer_comparator.cc:145] Difference at 129: 0, expected 0.406227
E0000 00:00:1746711197.089688 3123193 buffer_comparator.cc:145] Difference at 130: 0, expected 0.311948
E0000 00:00:1746711197.089691 3123193 buffer_comparator.cc:145] Difference at 131: 0, expected 0.53677
E0000 00:00:1746711197.089694 3123193 buffer_comparator.cc:145] Difference at 132: 0, expected 0.172814
E0000 00:00:1746711197.089697 3123193 buffer_comparator.cc:145] Difference at 133: 0, expected 0.314312
E0000 00:00:1746711197.089700 3123193 buffer_comparator.cc:145] Difference at 134: 0, expected 1.17027
E0000 00:00:1746711197.089703 3123193 buffer_comparator.cc:145] Difference at 135: 0, expected 1.05396
E0000 00:00:1746711197.089706 3123193 buffer_comparator.cc:145] Difference at 136: 0, expected 0.788122
E0000 00:00:1746711197.089709 3123193 buffer_comparator.cc:145] Difference at 137: 0, expected 0.232274
2025-05-08 13:33:17.089714: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.092121 3123193 buffer_comparator.cc:145] Difference at 128: 0, expected 1.00573
E0000 00:00:1746711197.092133 3123193 buffer_comparator.cc:145] Difference at 129: 0, expected 0.406227
E0000 00:00:1746711197.092137 3123193 buffer_comparator.cc:145] Difference at 130: 0, expected 0.311948
E0000 00:00:1746711197.092141 3123193 buffer_comparator.cc:145] Difference at 131: 0, expected 0.53677
E0000 00:00:1746711197.092144 3123193 buffer_comparator.cc:145] Difference at 132: 0, expected 0.172814
E0000 00:00:1746711197.092147 3123193 buffer_comparator.cc:145] Difference at 133: 0, expected 0.314312
E0000 00:00:1746711197.092150 3123193 buffer_comparator.cc:145] Difference at 134: 0, expected 1.17027
E0000 00:00:1746711197.092153 3123193 buffer_comparator.cc:145] Difference at 135: 0, expected 1.05396
E0000 00:00:1746711197.092156 3123193 buffer_comparator.cc:145] Difference at 136: 0, expected 0.788122
E0000 00:00:1746711197.092159 3123193 buffer_comparator.cc:145] Difference at 137: 0, expected 0.232274
2025-05-08 13:33:17.092164: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.094562 3123193 buffer_comparator.cc:145] Difference at 128: 0, expected 1.00573
E0000 00:00:1746711197.094574 3123193 buffer_comparator.cc:145] Difference at 129: 0, expected 0.406227
E0000 00:00:1746711197.094578 3123193 buffer_comparator.cc:145] Difference at 130: 0, expected 0.311948
E0000 00:00:1746711197.094581 3123193 buffer_comparator.cc:145] Difference at 131: 0, expected 0.53677
E0000 00:00:1746711197.094584 3123193 buffer_comparator.cc:145] Difference at 132: 0, expected 0.172814
E0000 00:00:1746711197.094587 3123193 buffer_comparator.cc:145] Difference at 133: 0, expected 0.314312
E0000 00:00:1746711197.094590 3123193 buffer_comparator.cc:145] Difference at 134: 0, expected 1.17027
E0000 00:00:1746711197.094593 3123193 buffer_comparator.cc:145] Difference at 135: 0, expected 1.05396
E0000 00:00:1746711197.094596 3123193 buffer_comparator.cc:145] Difference at 136: 0, expected 0.788122
E0000 00:00:1746711197.094599 3123193 buffer_comparator.cc:145] Difference at 137: 0, expected 0.232274
2025-05-08 13:33:17.094604: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.097001 3123193 buffer_comparator.cc:145] Difference at 128: 0, expected 1.00573
E0000 00:00:1746711197.097014 3123193 buffer_comparator.cc:145] Difference at 129: 0, expected 0.406227
E0000 00:00:1746711197.097018 3123193 buffer_comparator.cc:145] Difference at 130: 0, expected 0.311948
E0000 00:00:1746711197.097021 3123193 buffer_comparator.cc:145] Difference at 131: 0, expected 0.53677
E0000 00:00:1746711197.097024 3123193 buffer_comparator.cc:145] Difference at 132: 0, expected 0.172814
E0000 00:00:1746711197.097027 3123193 buffer_comparator.cc:145] Difference at 133: 0, expected 0.314312
E0000 00:00:1746711197.097030 3123193 buffer_comparator.cc:145] Difference at 134: 0, expected 1.17027
E0000 00:00:1746711197.097033 3123193 buffer_comparator.cc:145] Difference at 135: 0, expected 1.05396
E0000 00:00:1746711197.097036 3123193 buffer_comparator.cc:145] Difference at 136: 0, expected 0.788122
E0000 00:00:1746711197.097039 3123193 buffer_comparator.cc:145] Difference at 137: 0, expected 0.232274
2025-05-08 13:33:17.097043: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.099446 3123193 buffer_comparator.cc:145] Difference at 128: 0, expected 1.00573
E0000 00:00:1746711197.099458 3123193 buffer_comparator.cc:145] Difference at 129: 0, expected 0.406227
E0000 00:00:1746711197.099462 3123193 buffer_comparator.cc:145] Difference at 130: 0, expected 0.311948
E0000 00:00:1746711197.099465 3123193 buffer_comparator.cc:145] Difference at 131: 0, expected 0.53677
E0000 00:00:1746711197.099468 3123193 buffer_comparator.cc:145] Difference at 132: 0, expected 0.172814
E0000 00:00:1746711197.099471 3123193 buffer_comparator.cc:145] Difference at 133: 0, expected 0.314312
E0000 00:00:1746711197.099474 3123193 buffer_comparator.cc:145] Difference at 134: 0, expected 1.17027
E0000 00:00:1746711197.099478 3123193 buffer_comparator.cc:145] Difference at 135: 0, expected 1.05396
E0000 00:00:1746711197.099481 3123193 buffer_comparator.cc:145] Difference at 136: 0, expected 0.788122
E0000 00:00:1746711197.099484 3123193 buffer_comparator.cc:145] Difference at 137: 0, expected 0.232274
2025-05-08 13:33:17.099489: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711197.102185 3123193 buffer_comparator.cc:145] Difference at 32: -nan, expected 1.62244
E0000 00:00:1746711197.102197 3123193 buffer_comparator.cc:145] Difference at 33: -nan, expected 1.87084
E0000 00:00:1746711197.102201 3123193 buffer_comparator.cc:145] Difference at 34: -nan, expected 1.07351
E0000 00:00:1746711197.102203 3123193 buffer_comparator.cc:145] Difference at 35: -nan, expected 2.92445
E0000 00:00:1746711197.102206 3123193 buffer_comparator.cc:145] Difference at 36: -nan, expected 1.98056
E0000 00:00:1746711197.102209 3123193 buffer_comparator.cc:145] Difference at 37: -nan, expected 2.07715
E0000 00:00:1746711197.102212 3123193 buffer_comparator.cc:145] Difference at 38: -nan, expected 1.56458
E0000 00:00:1746711197.102214 3123193 buffer_comparator.cc:145] Difference at 39: -nan, expected 2.27034
E0000 00:00:1746711197.102217 3123193 buffer_comparator.cc:145] Difference at 40: -nan, expected 2.31795
E0000 00:00:1746711197.102220 3123193 buffer_comparator.cc:145] Difference at 41: -nan, expected 2.55731
2025-05-08 13:33:17.102224: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.777501 3123193 buffer_comparator.cc:145] Difference at 16: 0.196842, expected 9.68745
E0000 00:00:1746711234.777550 3123193 buffer_comparator.cc:145] Difference at 17: 0.688536, expected 10.1876
E0000 00:00:1746711234.777557 3123193 buffer_comparator.cc:145] Difference at 18: 0.927057, expected 8.84104
E0000 00:00:1746711234.777560 3123193 buffer_comparator.cc:145] Difference at 19: 0.579189, expected 10.0381
E0000 00:00:1746711234.777563 3123193 buffer_comparator.cc:145] Difference at 20: 0.374055, expected 7.30446
E0000 00:00:1746711234.777566 3123193 buffer_comparator.cc:145] Difference at 21: 0.216797, expected 8.26483
E0000 00:00:1746711234.777569 3123193 buffer_comparator.cc:145] Difference at 22: 0.731212, expected 10.8549
E0000 00:00:1746711234.777572 3123193 buffer_comparator.cc:145] Difference at 23: 0.700668, expected 7.87482
E0000 00:00:1746711234.777575 3123193 buffer_comparator.cc:145] Difference at 24: 0.5317, expected 9.78239
E0000 00:00:1746711234.777578 3123193 buffer_comparator.cc:145] Difference at 25: 0.24009, expected 11.3838
2025-05-08 13:33:54.777590: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.779652 3123193 buffer_comparator.cc:145] Difference at 16: 0.196842, expected 9.68745
E0000 00:00:1746711234.779665 3123193 buffer_comparator.cc:145] Difference at 17: 0.688536, expected 10.1876
E0000 00:00:1746711234.779669 3123193 buffer_comparator.cc:145] Difference at 18: 0.927057, expected 8.84104
E0000 00:00:1746711234.779672 3123193 buffer_comparator.cc:145] Difference at 19: 0.579189, expected 10.0381
E0000 00:00:1746711234.779675 3123193 buffer_comparator.cc:145] Difference at 20: 0.374055, expected 7.30446
E0000 00:00:1746711234.779678 3123193 buffer_comparator.cc:145] Difference at 21: 0.216797, expected 8.26483
E0000 00:00:1746711234.779681 3123193 buffer_comparator.cc:145] Difference at 22: 0.731212, expected 10.8549
E0000 00:00:1746711234.779684 3123193 buffer_comparator.cc:145] Difference at 23: 0.700668, expected 7.87482
E0000 00:00:1746711234.779687 3123193 buffer_comparator.cc:145] Difference at 24: 0.5317, expected 9.78239
E0000 00:00:1746711234.779690 3123193 buffer_comparator.cc:145] Difference at 25: 0.24009, expected 11.3838
2025-05-08 13:33:54.779697: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.781721 3123193 buffer_comparator.cc:145] Difference at 16: 0.196842, expected 9.68745
E0000 00:00:1746711234.781734 3123193 buffer_comparator.cc:145] Difference at 17: 0.688536, expected 10.1876
E0000 00:00:1746711234.781738 3123193 buffer_comparator.cc:145] Difference at 18: 0.927057, expected 8.84104
E0000 00:00:1746711234.781741 3123193 buffer_comparator.cc:145] Difference at 19: 0.579189, expected 10.0381
E0000 00:00:1746711234.781744 3123193 buffer_comparator.cc:145] Difference at 20: 0.374055, expected 7.30446
E0000 00:00:1746711234.781747 3123193 buffer_comparator.cc:145] Difference at 21: 0.216797, expected 8.26483
E0000 00:00:1746711234.781750 3123193 buffer_comparator.cc:145] Difference at 22: 0.731212, expected 10.8549
E0000 00:00:1746711234.781753 3123193 buffer_comparator.cc:145] Difference at 23: 0.700668, expected 7.87482
E0000 00:00:1746711234.781756 3123193 buffer_comparator.cc:145] Difference at 24: 0.5317, expected 9.78239
E0000 00:00:1746711234.781759 3123193 buffer_comparator.cc:145] Difference at 25: 0.24009, expected 11.3838
2025-05-08 13:33:54.781764: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.783776 3123193 buffer_comparator.cc:145] Difference at 32: 0.766695, expected 9.13848
E0000 00:00:1746711234.783788 3123193 buffer_comparator.cc:145] Difference at 33: 1.02114, expected 7.0792
E0000 00:00:1746711234.783791 3123193 buffer_comparator.cc:145] Difference at 34: 0.0917029, expected 10.2155
E0000 00:00:1746711234.783795 3123193 buffer_comparator.cc:145] Difference at 35: 0.842239, expected 9.45231
E0000 00:00:1746711234.783798 3123193 buffer_comparator.cc:145] Difference at 36: 0.52163, expected 10.5298
E0000 00:00:1746711234.783801 3123193 buffer_comparator.cc:145] Difference at 37: 0.313266, expected 9.84508
E0000 00:00:1746711234.783804 3123193 buffer_comparator.cc:145] Difference at 38: 1.04173, expected 9.51338
E0000 00:00:1746711234.783807 3123193 buffer_comparator.cc:145] Difference at 39: 0.974961, expected 10.1471
E0000 00:00:1746711234.783810 3123193 buffer_comparator.cc:145] Difference at 40: 0.978602, expected 9.57115
E0000 00:00:1746711234.783813 3123193 buffer_comparator.cc:145] Difference at 41: 1.00507, expected 8.63119
2025-05-08 13:33:54.783818: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.785864 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.785876 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.785880 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.785883 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.785887 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.785889 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.785892 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.785895 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.785898 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.785901 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.785906: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.788021 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.788033 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.788037 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.788040 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.788043 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.788046 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.788049 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.788052 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.788055 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.788058 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.788062: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.790162 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.790187 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.790191 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.790194 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.790197 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.790200 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.790203 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.790206 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.790208 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.790211 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.790217: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.792268 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.792280 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.792289 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.792292 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.792295 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.792298 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.792301 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.792304 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.792307 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.792310 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.792315: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.794352 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.794365 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.794370 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.794373 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.794376 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.794379 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.794382 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.794384 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.794387 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.794390 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.794395: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.796412 3123193 buffer_comparator.cc:145] Difference at 64: 0.586121, expected 9.67458
E0000 00:00:1746711234.796424 3123193 buffer_comparator.cc:145] Difference at 65: 0.809946, expected 10.734
E0000 00:00:1746711234.796428 3123193 buffer_comparator.cc:145] Difference at 66: 0.423876, expected 10.6109
E0000 00:00:1746711234.796431 3123193 buffer_comparator.cc:145] Difference at 67: 0.65869, expected 8.23326
E0000 00:00:1746711234.796434 3123193 buffer_comparator.cc:145] Difference at 68: 1.08471, expected 8.19665
E0000 00:00:1746711234.796437 3123193 buffer_comparator.cc:145] Difference at 69: 0.449177, expected 9.30282
E0000 00:00:1746711234.796440 3123193 buffer_comparator.cc:145] Difference at 70: 0.988388, expected 8.16784
E0000 00:00:1746711234.796443 3123193 buffer_comparator.cc:145] Difference at 71: 0.849879, expected 9.34399
E0000 00:00:1746711234.796446 3123193 buffer_comparator.cc:145] Difference at 72: 1.03034, expected 9.36502
E0000 00:00:1746711234.796449 3123193 buffer_comparator.cc:145] Difference at 73: 0.892119, expected 8.82565
2025-05-08 13:33:54.796454: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.801897 3123193 buffer_comparator.cc:145] Difference at 16: 9.68831, expected 34.2325
E0000 00:00:1746711234.801921 3123193 buffer_comparator.cc:145] Difference at 17: 10.1886, expected 32.4845
E0000 00:00:1746711234.801925 3123193 buffer_comparator.cc:145] Difference at 18: 8.84087, expected 35.8503
E0000 00:00:1746711234.801928 3123193 buffer_comparator.cc:145] Difference at 19: 10.0385, expected 38.0823
E0000 00:00:1746711234.801932 3123193 buffer_comparator.cc:145] Difference at 20: 7.30459, expected 32.6811
E0000 00:00:1746711234.801935 3123193 buffer_comparator.cc:145] Difference at 21: 8.26478, expected 37.818
E0000 00:00:1746711234.801938 3123193 buffer_comparator.cc:145] Difference at 22: 10.8556, expected 35.4896
E0000 00:00:1746711234.801941 3123193 buffer_comparator.cc:145] Difference at 23: 7.87467, expected 35.057
E0000 00:00:1746711234.801944 3123193 buffer_comparator.cc:145] Difference at 24: 9.78306, expected 37.6513
E0000 00:00:1746711234.801947 3123193 buffer_comparator.cc:145] Difference at 25: 11.3832, expected 36.0917
2025-05-08 13:33:54.801954: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.803961 3123193 buffer_comparator.cc:145] Difference at 16: 9.68831, expected 34.2325
E0000 00:00:1746711234.803973 3123193 buffer_comparator.cc:145] Difference at 17: 10.1886, expected 32.4845
E0000 00:00:1746711234.803977 3123193 buffer_comparator.cc:145] Difference at 18: 8.84087, expected 35.8503
E0000 00:00:1746711234.803982 3123193 buffer_comparator.cc:145] Difference at 19: 10.0385, expected 38.0823
E0000 00:00:1746711234.803985 3123193 buffer_comparator.cc:145] Difference at 20: 7.30459, expected 32.6811
E0000 00:00:1746711234.803988 3123193 buffer_comparator.cc:145] Difference at 21: 8.26478, expected 37.818
E0000 00:00:1746711234.803991 3123193 buffer_comparator.cc:145] Difference at 22: 10.8556, expected 35.4896
E0000 00:00:1746711234.803995 3123193 buffer_comparator.cc:145] Difference at 23: 7.87467, expected 35.057
E0000 00:00:1746711234.803998 3123193 buffer_comparator.cc:145] Difference at 24: 9.78306, expected 37.6513
E0000 00:00:1746711234.804001 3123193 buffer_comparator.cc:145] Difference at 25: 11.3832, expected 36.0917
2025-05-08 13:33:54.804005: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.806001 3123193 buffer_comparator.cc:145] Difference at 16: 9.68831, expected 34.2325
E0000 00:00:1746711234.806013 3123193 buffer_comparator.cc:145] Difference at 17: 10.1886, expected 32.4845
E0000 00:00:1746711234.806017 3123193 buffer_comparator.cc:145] Difference at 18: 8.84087, expected 35.8503
E0000 00:00:1746711234.806021 3123193 buffer_comparator.cc:145] Difference at 19: 10.0385, expected 38.0823
E0000 00:00:1746711234.806024 3123193 buffer_comparator.cc:145] Difference at 20: 7.30459, expected 32.6811
E0000 00:00:1746711234.806027 3123193 buffer_comparator.cc:145] Difference at 21: 8.26478, expected 37.818
E0000 00:00:1746711234.806030 3123193 buffer_comparator.cc:145] Difference at 22: 10.8556, expected 35.4896
E0000 00:00:1746711234.806033 3123193 buffer_comparator.cc:145] Difference at 23: 7.87467, expected 35.057
E0000 00:00:1746711234.806036 3123193 buffer_comparator.cc:145] Difference at 24: 9.78306, expected 37.6513
E0000 00:00:1746711234.806039 3123193 buffer_comparator.cc:145] Difference at 25: 11.3832, expected 36.0917
2025-05-08 13:33:54.806044: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.808040 3123193 buffer_comparator.cc:145] Difference at 16: 9.68831, expected 34.2325
E0000 00:00:1746711234.808052 3123193 buffer_comparator.cc:145] Difference at 17: 10.1886, expected 32.4845
E0000 00:00:1746711234.808057 3123193 buffer_comparator.cc:145] Difference at 18: 8.84087, expected 35.8503
E0000 00:00:1746711234.808060 3123193 buffer_comparator.cc:145] Difference at 19: 10.0385, expected 38.0823
E0000 00:00:1746711234.808063 3123193 buffer_comparator.cc:145] Difference at 20: 7.30459, expected 32.6811
E0000 00:00:1746711234.808066 3123193 buffer_comparator.cc:145] Difference at 21: 8.26478, expected 37.818
E0000 00:00:1746711234.808069 3123193 buffer_comparator.cc:145] Difference at 22: 10.8556, expected 35.4896
E0000 00:00:1746711234.808072 3123193 buffer_comparator.cc:145] Difference at 23: 7.87467, expected 35.057
E0000 00:00:1746711234.808075 3123193 buffer_comparator.cc:145] Difference at 24: 9.78306, expected 37.6513
E0000 00:00:1746711234.808078 3123193 buffer_comparator.cc:145] Difference at 25: 11.3832, expected 36.0917
2025-05-08 13:33:54.808083: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.810090 3123193 buffer_comparator.cc:145] Difference at 16: 9.68831, expected 34.2325
E0000 00:00:1746711234.810102 3123193 buffer_comparator.cc:145] Difference at 17: 10.1886, expected 32.4845
E0000 00:00:1746711234.810106 3123193 buffer_comparator.cc:145] Difference at 18: 8.84087, expected 35.8503
E0000 00:00:1746711234.810110 3123193 buffer_comparator.cc:145] Difference at 19: 10.0385, expected 38.0823
E0000 00:00:1746711234.810113 3123193 buffer_comparator.cc:145] Difference at 20: 7.30459, expected 32.6811
E0000 00:00:1746711234.810117 3123193 buffer_comparator.cc:145] Difference at 21: 8.26478, expected 37.818
E0000 00:00:1746711234.810120 3123193 buffer_comparator.cc:145] Difference at 22: 10.8556, expected 35.4896
E0000 00:00:1746711234.810123 3123193 buffer_comparator.cc:145] Difference at 23: 7.87467, expected 35.057
E0000 00:00:1746711234.810126 3123193 buffer_comparator.cc:145] Difference at 24: 9.78306, expected 37.6513
E0000 00:00:1746711234.810129 3123193 buffer_comparator.cc:145] Difference at 25: 11.3832, expected 36.0917
2025-05-08 13:33:54.810134: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.830726 3123193 buffer_comparator.cc:145] Difference at 16: -nan, expected 34.687
E0000 00:00:1746711234.830758 3123193 buffer_comparator.cc:145] Difference at 17: -nan, expected 32.6585
E0000 00:00:1746711234.830765 3123193 buffer_comparator.cc:145] Difference at 18: -nan, expected 37.2083
E0000 00:00:1746711234.830768 3123193 buffer_comparator.cc:145] Difference at 19: -nan, expected 32.2063
E0000 00:00:1746711234.830771 3123193 buffer_comparator.cc:145] Difference at 20: -nan, expected 33.4727
E0000 00:00:1746711234.830774 3123193 buffer_comparator.cc:145] Difference at 21: -nan, expected 33.0033
E0000 00:00:1746711234.830777 3123193 buffer_comparator.cc:145] Difference at 22: -nan, expected 31.6193
E0000 00:00:1746711234.830780 3123193 buffer_comparator.cc:145] Difference at 23: -nan, expected 32.1492
E0000 00:00:1746711234.830783 3123193 buffer_comparator.cc:145] Difference at 24: -nan, expected 32.5713
E0000 00:00:1746711234.830785 3123193 buffer_comparator.cc:145] Difference at 25: -nan, expected 36.4575
2025-05-08 13:33:54.830793: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.832802 3123193 buffer_comparator.cc:145] Difference at 16: -nan, expected 34.687
E0000 00:00:1746711234.832814 3123193 buffer_comparator.cc:145] Difference at 17: -nan, expected 32.6585
E0000 00:00:1746711234.832818 3123193 buffer_comparator.cc:145] Difference at 18: -nan, expected 37.2083
E0000 00:00:1746711234.832821 3123193 buffer_comparator.cc:145] Difference at 19: -nan, expected 32.2063
E0000 00:00:1746711234.832824 3123193 buffer_comparator.cc:145] Difference at 20: -nan, expected 33.4727
E0000 00:00:1746711234.832827 3123193 buffer_comparator.cc:145] Difference at 21: -nan, expected 33.0033
E0000 00:00:1746711234.832830 3123193 buffer_comparator.cc:145] Difference at 22: -nan, expected 31.6193
E0000 00:00:1746711234.832832 3123193 buffer_comparator.cc:145] Difference at 23: -nan, expected 32.1492
E0000 00:00:1746711234.832835 3123193 buffer_comparator.cc:145] Difference at 24: -nan, expected 32.5713
E0000 00:00:1746711234.832838 3123193 buffer_comparator.cc:145] Difference at 25: -nan, expected 36.4575
2025-05-08 13:33:54.832843: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.834857 3123193 buffer_comparator.cc:145] Difference at 16: -nan, expected 34.687
E0000 00:00:1746711234.834869 3123193 buffer_comparator.cc:145] Difference at 17: -nan, expected 32.6585
E0000 00:00:1746711234.834873 3123193 buffer_comparator.cc:145] Difference at 18: -nan, expected 37.2083
E0000 00:00:1746711234.834876 3123193 buffer_comparator.cc:145] Difference at 19: -nan, expected 32.2063
E0000 00:00:1746711234.834879 3123193 buffer_comparator.cc:145] Difference at 20: -nan, expected 33.4727
E0000 00:00:1746711234.834882 3123193 buffer_comparator.cc:145] Difference at 21: -nan, expected 33.0033
E0000 00:00:1746711234.834885 3123193 buffer_comparator.cc:145] Difference at 22: -nan, expected 31.6193
E0000 00:00:1746711234.834888 3123193 buffer_comparator.cc:145] Difference at 23: -nan, expected 32.1492
E0000 00:00:1746711234.834892 3123193 buffer_comparator.cc:145] Difference at 24: -nan, expected 32.5713
E0000 00:00:1746711234.834895 3123193 buffer_comparator.cc:145] Difference at 25: -nan, expected 36.4575
2025-05-08 13:33:54.834900: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.836895 3123193 buffer_comparator.cc:145] Difference at 16: -nan, expected 34.687
E0000 00:00:1746711234.836907 3123193 buffer_comparator.cc:145] Difference at 17: -nan, expected 32.6585
E0000 00:00:1746711234.836910 3123193 buffer_comparator.cc:145] Difference at 18: -nan, expected 37.2083
E0000 00:00:1746711234.836913 3123193 buffer_comparator.cc:145] Difference at 19: -nan, expected 32.2063
E0000 00:00:1746711234.836916 3123193 buffer_comparator.cc:145] Difference at 20: -nan, expected 33.4727
E0000 00:00:1746711234.836919 3123193 buffer_comparator.cc:145] Difference at 21: -nan, expected 33.0033
E0000 00:00:1746711234.836922 3123193 buffer_comparator.cc:145] Difference at 22: -nan, expected 31.6193
E0000 00:00:1746711234.836925 3123193 buffer_comparator.cc:145] Difference at 23: -nan, expected 32.1492
E0000 00:00:1746711234.836927 3123193 buffer_comparator.cc:145] Difference at 24: -nan, expected 32.5713
E0000 00:00:1746711234.836930 3123193 buffer_comparator.cc:145] Difference at 25: -nan, expected 36.4575
2025-05-08 13:33:54.836935: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1746711234.838930 3123193 buffer_comparator.cc:145] Difference at 16: -nan, expected 34.687
E0000 00:00:1746711234.838944 3123193 buffer_comparator.cc:145] Difference at 17: -nan, expected 32.6585
E0000 00:00:1746711234.838948 3123193 buffer_comparator.cc:145] Difference at 18: -nan, expected 37.2083
E0000 00:00:1746711234.838952 3123193 buffer_comparator.cc:145] Difference at 19: -nan, expected 32.2063
E0000 00:00:1746711234.838955 3123193 buffer_comparator.cc:145] Difference at 20: -nan, expected 33.4727
E0000 00:00:1746711234.838958 3123193 buffer_comparator.cc:145] Difference at 21: -nan, expected 33.0033
E0000 00:00:1746711234.838961 3123193 buffer_comparator.cc:145] Difference at 22: -nan, expected 31.6193
E0000 00:00:1746711234.838963 3123193 buffer_comparator.cc:145] Difference at 23: -nan, expected 32.1492
E0000 00:00:1746711234.838966 3123193 buffer_comparator.cc:145] Difference at 24: -nan, expected 32.5713
E0000 00:00:1746711234.838969 3123193 buffer_comparator.cc:145] Difference at 25: -nan, expected 36.4575
2025-05-08 13:33:54.838974: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1179] Results do not match the reference. This is likely a bug/unexpected loss of precision.
Epoch [ 1]: Loss 0.77102
Validation: Loss 0.71518 Accuracy 0.49219
Epoch [ 2]: Loss 0.66183
Validation: Loss 0.58588 Accuracy 0.49219
Epoch [ 3]: Loss 0.52172
Validation: Loss 0.43412 Accuracy 1.00000
Epoch [ 4]: Loss 0.37535
Validation: Loss 0.30415 Accuracy 1.00000
Epoch [ 5]: Loss 0.26135
Validation: Loss 0.21227 Accuracy 1.00000
Epoch [ 6]: Loss 0.18548
Validation: Loss 0.15504 Accuracy 1.00000
Epoch [ 7]: Loss 0.13831
Validation: Loss 0.11817 Accuracy 1.00000
Epoch [ 8]: Loss 0.10739
Validation: Loss 0.09384 Accuracy 1.00000
Epoch [ 9]: Loss 0.08644
Validation: Loss 0.07665 Accuracy 1.00000
Epoch [ 10]: Loss 0.07124
Validation: Loss 0.06392 Accuracy 1.00000
Epoch [ 11]: Loss 0.05986
Validation: Loss 0.05431 Accuracy 1.00000
Epoch [ 12]: Loss 0.05111
Validation: Loss 0.04649 Accuracy 1.00000
Epoch [ 13]: Loss 0.04373
Validation: Loss 0.03968 Accuracy 1.00000
Epoch [ 14]: Loss 0.03703
Validation: Loss 0.03322 Accuracy 1.00000
Epoch [ 15]: Loss 0.03061
Validation: Loss 0.02685 Accuracy 1.00000
Epoch [ 16]: Loss 0.02452
Validation: Loss 0.02131 Accuracy 1.00000
Epoch [ 17]: Loss 0.01978
Validation: Loss 0.01757 Accuracy 1.00000
Epoch [ 18]: Loss 0.01659
Validation: Loss 0.01510 Accuracy 1.00000
Epoch [ 19]: Loss 0.01432
Validation: Loss 0.01322 Accuracy 1.00000
Epoch [ 20]: Loss 0.01265
Validation: Loss 0.01163 Accuracy 1.00000
Epoch [ 21]: Loss 0.01110
Validation: Loss 0.01018 Accuracy 1.00000
Epoch [ 22]: Loss 0.00971
Validation: Loss 0.00878 Accuracy 1.00000
Epoch [ 23]: Loss 0.00831
Validation: Loss 0.00747 Accuracy 1.00000
Epoch [ 24]: Loss 0.00706
Validation: Loss 0.00638 Accuracy 1.00000
Epoch [ 25]: Loss 0.00610
Validation: Loss 0.00559 Accuracy 1.00000We can also train the compact model with the exact same code!
julia
ps_trained2, st_trained2 = main(SpiralClassifierCompact)┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
└ @ LuxCore /var/lib/buildkite-agent/builds/gpuci-10/julialang/lux-dot-jl/lib/LuxCore/src/LuxCore.jl:18
Epoch [ 1]: Loss 0.94544
Validation: Loss 0.69613 Accuracy 0.56250
Epoch [ 2]: Loss 0.61169
Validation: Loss 0.50845 Accuracy 1.00000
Epoch [ 3]: Loss 0.45723
Validation: Loss 0.43266 Accuracy 1.00000
Epoch [ 4]: Loss 0.38456
Validation: Loss 0.38257 Accuracy 1.00000
Epoch [ 5]: Loss 0.33964
Validation: Loss 0.33890 Accuracy 1.00000
Epoch [ 6]: Loss 0.29663
Validation: Loss 0.29796 Accuracy 1.00000
Epoch [ 7]: Loss 0.25725
Validation: Loss 0.26062 Accuracy 1.00000
Epoch [ 8]: Loss 0.22534
Validation: Loss 0.22709 Accuracy 1.00000
Epoch [ 9]: Loss 0.19821
Validation: Loss 0.19755 Accuracy 1.00000
Epoch [ 10]: Loss 0.17055
Validation: Loss 0.17127 Accuracy 1.00000
Epoch [ 11]: Loss 0.14845
Validation: Loss 0.14720 Accuracy 1.00000
Epoch [ 12]: Loss 0.12852
Validation: Loss 0.12489 Accuracy 1.00000
Epoch [ 13]: Loss 0.10823
Validation: Loss 0.10390 Accuracy 1.00000
Epoch [ 14]: Loss 0.08985
Validation: Loss 0.08369 Accuracy 1.00000
Epoch [ 15]: Loss 0.07280
Validation: Loss 0.06836 Accuracy 1.00000
Epoch [ 16]: Loss 0.06158
Validation: Loss 0.05989 Accuracy 1.00000
Epoch [ 17]: Loss 0.05491
Validation: Loss 0.05404 Accuracy 1.00000
Epoch [ 18]: Loss 0.04983
Validation: Loss 0.04915 Accuracy 1.00000
Epoch [ 19]: Loss 0.04534
Validation: Loss 0.04498 Accuracy 1.00000
Epoch [ 20]: Loss 0.04150
Validation: Loss 0.04140 Accuracy 1.00000
Epoch [ 21]: Loss 0.03821
Validation: Loss 0.03831 Accuracy 1.00000
Epoch [ 22]: Loss 0.03541
Validation: Loss 0.03558 Accuracy 1.00000
Epoch [ 23]: Loss 0.03282
Validation: Loss 0.03319 Accuracy 1.00000
Epoch [ 24]: Loss 0.03068
Validation: Loss 0.03105 Accuracy 1.00000
Epoch [ 25]: Loss 0.02865
Validation: Loss 0.02915 Accuracy 1.00000Saving the Model
We can save the model using JLD2 (and any other serialization library of your choice) Note that we transfer the model to CPU before saving. Additionally, we recommend that you don't save the model struct and only save the parameters and states.
julia
@save "trained_model.jld2" ps_trained st_trainedLet's try loading the model
julia
@load "trained_model.jld2" ps_trained st_trained2-element Vector{Symbol}:
:ps_trained
:st_trainedAppendix
julia
using InteractiveUtils
InteractiveUtils.versioninfo()
if @isdefined(MLDataDevices)
if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
println()
CUDA.versioninfo()
end
if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
println()
AMDGPU.versioninfo()
end
endJulia Version 1.11.5
Commit 760b2e5b739 (2025-04-14 06:53 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 48 × AMD EPYC 7402 24-Core Processor
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, znver2)
Threads: 48 default, 0 interactive, 24 GC (on 2 virtual cores)
Environment:
JULIA_CPU_THREADS = 2
LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
JULIA_PKG_SERVER =
JULIA_NUM_THREADS = 48
JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
JULIA_PKG_PRECOMPILE_AUTO = 0
JULIA_DEBUG = Literate
JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6This page was generated using Literate.jl.