Google Colab GPU Error

[rohaldb]

When running es_clip in the accompanying google colab instance with more than 0 GPU's, I get the following error:

Traceback (most recent call last):
  File "/usr/lib/python3.7/multiprocessing/process.py", line 297, in _bootstrap
    self.run()
  File "/usr/lib/python3.7/multiprocessing/process.py", line 99, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/lib/python3.7/multiprocessing/pool.py", line 105, in worker
    initializer(*initargs)
  File "/content/brain-tokyo-workshop/es-clip/es_clip.py", line 97, in init_worker
    model, preprocess = clip.load('ViT-B/32', device, jit=True)  # change jit=True
  File "/usr/local/lib/python3.7/dist-packages/clip/clip.py", line 137, in load
    model = build_model(state_dict or model.state_dict()).to(device)
  File "/usr/local/lib/python3.7/dist-packages/clip/model.py", line 396, in build_model
    vit = "visual.proj" in state_dict
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py", line 816, in __contains__
    return self.forward_magic_method("__contains__", key)
  File "/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py", line 803, in forward_magic_method
    raise NotImplementedError()

Has anyone figured out how to solve this?

[alantian]

Hey @rohaldb , thanks for your interest on ES-CLIP!

We have re-implemented ES-CLIP in JAX, which is more suitable for running on Colab here:

• https://github.com/google/evojax/blob/main/examples/notebooks/AbstractPainting01.ipynb
• https://github.com/google/evojax/blob/main/examples/notebooks/AbstractPainting02.ipynb

Maybe you want to check it out?

[rohaldb]

Ah excellent, thanks!