Torch - 随机选择要切片的张量维度。

我有一个大小为 (1 x n x n x n) 的 torch tensor，我想随机选择最后三个维度中的一个，以 s 为标准随机切片，然后处理。例如，它可以以相等的概率输出以下张量;

(1 x s x n x n)
(1 x n x s x n)
(1 x n x n x s)

我意识到我可以使用一些 if else 语句来实现，但我想知道是否有一个更“简洁”的选择，使用类似 torch.random(1, 4) 的函数来选择维度。