HPUStrategy Implementation Error

I am running a benchmark on the AWS Habana Base AMI, Ubuntu 18.04. I’ve cloned my github repo and I am running my test in the TF 2.7.0 docker image. Running on a single HPU works fine but when trying to implement HPUStrategy for multi-worker training I run into problems. When creating the strategy using HPUStrategy, I keep getting this error:

TypeError: _initialize_local() got an unexpected keyword argument ‘device’

Is there an additional step I need to take before creating the strategy? I’ve added a screenshot of the script and the error occurs on line 11.


Hi @izzy, thanks for posting your question. I’m assuming you are using the Habana Base AMI based on the 1.2.0 SynapseAI Software release and the associated TF docker image? Please confirm.

Can you please give some more context where you are getting this error?

Your basic code block looks good, so there’s not too much to go on here, but you can take a look at our tf.distribute examples on our GitHub page here