pytorchの自動微分(backward)のエラー回避方法

前提・実現したいこと

pytorchのbackward機能を使って、とある場所のみの勾配を得たく

pytorch
1self.features.backward(gradient=one_hot, retain_graph=True)

とone_hotのtensorでマスクしようとしています。

発生している問題・エラーメッセージ

backward()を走ったところで下のエラーが出ます。

RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED

該当のソースコード

pytorch
1    def _encode_one_hot(self, ids):
2        # self.features.size() torch.Size([1, 20, 96, 96])
3        one_hot = torch.zeros_like(self.features).to(self.device)
4        one_hot[ids] = 1
5        return one_hot
6
7    def _backward(self, ids):
8        one_hot = self._encode_one_hot(ids)
9        self.model.zero_grad()
10        self.features.backward(gradient=one_hot, retain_graph=True)
11
12    ids = torch.tensor([0, 11, 1, 1])
13    self._backward(ids=ids)

試したこと

one_hotにしたい配列のインデックスを直打ちするとエラーが出ないのですが、動的にインデックスを変更する必要があるため、引数でone_hotにする位置を指定できるようにしたいです。

pytorch
1    def _encode_one_hot(self, ids):
2        # self.features.size() torch.Size([1, 20, 96, 96])
3        one_hot = torch.zeros_like(self.features).to(self.device)
4        one_hot[0, 11, 1, 1] = 1　# ← インデックスを直打ちに変更
5        return one_hot