Compile Error due to a type mismatch

[Mostafa86]

I am getting the following error

(CompileError) deps/axon/lib/axon/loop.ex:469: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.

While trying to enforce the following policy

policy = Axon.MixedPrecision.create_policy( params: {:f, 64}, compute: {:f, 64}, output: {:f, 64} )

The error seems to be due to the gradient_state below initiated as :f32

axon/lib/axon/loop.ex

Line 360 in ddc49cc

gradient_state: zeros_like(model_state, type: :f32),

[seanmor5]

Ah yeah the gradient state needs to match the policy. I will look into this. In the mean time is there a reason you need f64 specifically?

[Mostafa86]

@seanmor5 Thanks 👍

Regarding f64, not really it is just an error that I came across while trying to hunt down the reason behind a NaN in one of the models.