AbstractPainting02.ipynb. doesn't work on colab

[shi3z]

Hello, this is a really great code.

I was able to run "Abstract Painting 01" very well at Google coab.
However, when I ran "AbstractPainting02", an error occurred.

Exception                                 Traceback (most recent call last)
[<ipython-input-20-b16203d22159>](https://localhost:8080/#) in <module>()
      2 devices = jax.local_devices()
      3 
----> 4 image_fn, text_fn, jax_params, jax_preprocess = clip_jax.load('ViT-B/32', "cpu")
      5 
      6 target_text_ids = jnp.array(clip_jax.tokenize([prompt])) # already with batch dim

3 frames
[/content/CLIP_JAX/clip_jax/clip.py](https://localhost:8080/#) in process_node(value, name)
    117             new_tensor = jnp.array(pytorch_tensor)
    118         else:
--> 119             raise Exception("not implemented")
    120 
    121         assert new_tensor.shape == value.shape

Exception: not implemented

Which version of clip_jax when you made?

Best

[hardmaru]

@alantian

[alantian]

Hey @shi3z thanks for reporting this issue!

We can reproduce the issue and have identified it's related to a recent change in haiku (https://github.com/deepmind/dm-haiku/releases/tag/v0.0.6 , which was released 2 days ago).

Pushed a quick fix to our notebooks solving it for now.

[alantian]

@shi3z May I know if it works now on your side?

[shi3z]

@alantian I'm sorry my reply is late. Thanks to you, I was able to confirm that it works well. I'm excited to encounter amazing software!

Best