init_model.py

from octo.model.octo_model import OctoModel
import os

# Set JAX to CPU (consistent with your setup)
os.environ['JAX_PLATFORMS'] = 'cpu'

# Load the model to cache it
model = OctoModel.load_pretrained("hf://rail-berkeley/octo-small-1.5")
print("Model downloaded and cached successfully.")