how to change this torch to flax?

Question

lu-ming-lei opened this issue 3 years ago · comments

optimizer = torch.optim.Adam([variable1, variable2], 0.001)

I try this but no success: optimizer = flax.optim.Adam([variable1, variable2], 0.001)

joao guilherme · Answer 1 · Wed Jun 23 2021 03:30:29 GMT+0800 (China Standard Time)

Hey @lu-ming-lei can you provide a longer code snippet? In Flax we don't pass the variables we want to optmize when initialising the optimizer

lu-ming-lei · Answer 2 · Thu Jun 24 2021 22:51:54 GMT+0800 (China Standard Time)

optimizer = torch.optim.Adam([
  {'params': variables1, 'lr':0.1},
   {'params': variables2, 'lr':0.01},
  ])
how to change this code into flax?