Make Python interface work for non-SM80 targets (#726)

* Make Python interface work for non-SM80 targets

* Remove line in README
This commit is contained in:
Jack Kosaian
2022-12-07 21:53:33 -05:00
committed by GitHub
parent d6117ca362
commit df81d847d7
33 changed files with 149 additions and 20 deletions
@@ -34,6 +34,7 @@ import pycutlass
from pycutlass import *
from pycutlass.conv2d_operation import *
from pycutlass.utils import reference_model
from pycutlass.utils.device import device_cc
import sys
import torch.nn.functional as F
@@ -146,6 +147,11 @@ try:
except:
sys.exit(0)
cc = device_cc()
if args.compute_capability != cc:
raise Exception(("Parameter --compute-capability of {} "
"does not match that of the device of {}.").format(args.compute_capability, cc))
pycutlass.get_memory_pool(init_pool_size=2**30, max_pool_size=2**32)
np.random.seed(0)