Closed EloyAnguiano closed 5 months ago
Thanks for your message, this is totally expected in Haiku, because as the error says all hk.Module
s must be initialized inside an hk.transform
, and GraphCast is a Haiku module.
See this bit of code in the GraphCast demo:
@hk.transform_with_state
def run_forward(model_config, task_config, inputs, targets_template, forcings):
predictor = construct_wrapped_graphcast(model_config, task_config)
return predictor(inputs, targets_template=targets_template, forcings=forcings)
which contains an example of how to use hk.transform
.
Hi. Ia am trying to execute the graphcast model in a conda enviornment built with the same packages version of a working execution at google collab but whenever I try to build the model at
construct_wrapped_graphcast
function returns this error:I checked that both dm-haiku versions (collaboratory and local) are
0.0.11
. Is there any dockerfile to build a working environment or something like that? It is very difficult to run the same collab env at local.How to reproduce: