Cuda OOM with same script as MPS on 3070

[Sebo-the-tramp]

Description:

Using the same code for MPS, the CUDA implementation keeps increasing the memory without never freeing it "supposedly". Everytime the densification branch is executed, the memory increments as shown in photo.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[9], [line 26](vscode-notebook-cell:?execution_count=9&line=26)
     [22](vscode-notebook-cell:?execution_count=9&line=22) optimizer.zero_grad()        
     [24](vscode-notebook-cell:?execution_count=9&line=24) loss.backward()
---> [26](vscode-notebook-cell:?execution_count=9&line=26) optimizer.step()
     [28](vscode-notebook-cell:?execution_count=9&line=28) loss_history.append(loss.item())
     [30](vscode-notebook-cell:?execution_count=9&line=30) if epoch % densification_interval == 0 and epoch > 0:

File ~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:34, in Optimizer.step(self)
     [30](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:30) def step(self):
     [31](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:31)   """
     [32](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:32)   Performs a single optimization step.
     [33](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:33)   """
---> [34](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/nn/optim.py:34)   Tensor.realize(*self.schedule_step())

File ~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:3522, in _metadata_wrapper.<locals>._wrapper(*args, **kwargs)
   [3521](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:3521) def _wrapper(*args, **kwargs):
-> [3522](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:3522)   if _METADATA.get() is not None: return fn(*args, **kwargs)
   [3524](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:3524)   if TRACEMETA >= 2:
   [3525](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:3525)     caller_frame = sys._getframe(frame := 1)

File ~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:225, in Tensor.realize(self, do_update_stats, *lst)
    [223](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:223) def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
    [224](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/tensor.py:224)   """Triggers the computation needed to create these Tensor(s)."""
...
File ~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/runtime/ops_gpu.py:14, in check(status)
     [13](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/runtime/ops_gpu.py:13) def check(status):
---> [14](https://file+.vscode-resource.vscode-cdn.net/home/lab/Documents/scsv/tinysplat/tinysplat_2D/~/micromamba/envs/tinysplat/lib/python3.13/site-packages/tinygrad/runtime/ops_gpu.py:14)   if status != 0: raise RuntimeError(f"OpenCL Error {status}: {cl_errors.get(status, 'Unknown error')}")

RuntimeError: OpenCL Error -4: CL_MEM_OBJECT_ALLOCATION_FAILURE

[Sebo-the-tramp]

W_values[start_index:end_index, :] = W_values[distinct_indices_tensor[:max_len_distinct_integers], :]

It boils down to this part of the code