allow for install of cpu torch on linux, relax torch upper bound

[rbavery]

Opening an issue to see if y'all are amenable to some toml changes that I think will allow for this to be installed on more setups.

In the toml it declares

[tool.uv.sources]
claymodel = { git = "https://github.com/Clay-foundation/model.git" }
geobench = { git = "https://github.com/ServiceNow/geo-bench.git" }
torch = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]

Specifying a platform marker like index = "pytorch-cpu", marker = "sys_platform != 'linux'" means we can't install cpu only torch with olmoearth_pretrain on linux platforms.

We use cpu only torch in our CI and would like to be able to do this. In our own packaging we instead declare a uv extra with each index. One would opt in to installing either the CPU or GPU version of torch like so, without platform restrictions

uv sync --extra torch-cpu
uv sync --extra torch-cu128

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "torch-cpu" },
  { index = "pytorch-cu128", extra = "torch-cu128" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "torch-cpu" },
  { index = "pytorch-cu128", extra = "torch-cu128" },
]
pytorch-triton = [ { index = "pytorch-cu128", extra = "torch-cu128" } ]

optional-dependencies.torch-cpu = [
  "torch>=2.8,<2.9",
  "torchvision>=0.23,<1",
]
optional-dependencies.torch-cu128 = [
  "pytorch-triton",
  "torch>=2.8,<2.9",
  "torchvision>=0.23,<1",
]

I think making this explicit can be helpful since a user could have reasons to use the cpu or gpu version if they are on linux.

Also could the torch upper bound be relaxed since 2.9 is out?

I made a PR that addresses these points if the maintainers are open to reviewing it! #440

I confirmed that most tests pass locally (still running) and both envs can be solved with all groups.

[Volgat]

@rbavery I think this PR brings a significant improvement because: 1. It makes installation more flexible 2. It follows Python packaging best practices 3. It's backward compatible Documenting the extras in the README would be the next natural step.