python 高速に大容量のデータの保存したい

機械学習でデータセットが大きすぎてエラーが出ます

機械学習の種類は強化学習です　そのせいでデータセットがメモリに保存される
という状態になってます

ツール　pytorch
現在の保存関連の操作
(入力(3,128,128)　行動(2)　報酬)を毎step追加し
データセット全部をメモリ .to("cpu")
学習用のバッチをgpu .to("cuda:0)
一定期間ごとに余分な部分を削除

エラーの原因はデータセットをメモリに保存しているためで
メモリではなくssdに保存すればこのエラーを回避できるはずです
ですがpickleやpandas では高次元に対応できなかったり　上書きできなかったりでいいものが見つかりません　
なにかいい保存方法はありませんか？知っている人教えてください

問題のデータセット関連のコード

class Memory:#memory__
    def __init__(self, max_size=1000):
        self.buffer = [deque() for _ in range(5)]
        self.max_size=(max_size*main.R2D2trin)+main.multireward_steps+main.R2D2trin
        self.recarent_memory=deque([[] for _ in range(5)])
    
    def cat_buffer(self):#大きすぎる部分をカット
        if self.max_size < len(self.buffer[4]):
            while self.max_size < len(self.buffer[4]):
                self.buffer[0].popleft()
                self.buffer[1].popleft()
                self.buffer[2].popleft()
                self.buffer[3].popleft()
                self.buffer[4].popleft()
    
    def add(self, experience):#追加
        for _,i in enumerate(experience):
            self.buffer[_].append(i)