pytorch/torch/_C/_cudnn.pyi
cyyever 23eb0a3201 Improve typing in torch/types.py (#145237)
Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145237
Approved by: https://github.com/XuehaiPan, https://github.com/albanD

Co-authored-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
2025-01-28 05:29:12 +00:00

15 lines
317 B
Python

from enum import IntEnum
# Defined in torch/csrc/cuda/shared/cudnn.cpp
is_cuda: bool
def getRuntimeVersion() -> tuple[int, int, int]: ...
def getCompileVersion() -> tuple[int, int, int]: ...
def getVersionInt() -> int: ...
class RNNMode(IntEnum):
rnn_relu = ...
rnn_tanh = ...
lstm = ...
gru = ...