mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
10c4b98ade
Summary: * Deletes all weak script decorators / associated data structures / methods * In order to keep supporting the standard library in script, this enables recursive script on any function defined in `torch.nn` * Most changes in `torch/nn` are the result of `ag -Q "weak" torch/nn/ -l | xargs sed -i '/weak/d'`, only `rnn.py` needed manual editing to use the `ignore` and `export` to continue supporting the overloaded `forward` methods * `Sequential`/`ModuleList` no longer need to be added to constants since they are compiled on demand This should also fix https://github.com/pytorch/pytorch/issues/22212 Pull Request resolved: https://github.com/pytorch/pytorch/pull/22212 Differential Revision: D15988346 Pulled By: driazati fbshipit-source-id: af223e3ad0580be895377312949997a70e988e4f
13 lines
417 B
Python
13 lines
417 B
Python
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
|
|
|
|
def affine_grid_generator(theta, size):
|
|
# type: (Tensor, List[int]) -> Tensor
|
|
if theta.is_cuda and cudnn.enabled and cudnn.is_acceptable(theta) and len(size) == 4 and size[0] < 65536:
|
|
N, C, H, W = size
|
|
ret = torch.cudnn_affine_grid_generator(theta, N, C, H, W)
|
|
else:
|
|
ret = torch.affine_grid_generator(theta, size)
|
|
return ret
|