[examples] added skeleton for pytorch wrapper

This commit is contained in:
Philippe Tillet
2019-05-03 14:30:06 -04:00
parent 208d1525de
commit 0d694445e6
8 changed files with 181 additions and 15 deletions

View File

@@ -0,0 +1,11 @@
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.cpp_extension import load
from torch.distributions import categorical
from itertools import product
conv_triton = load( 'conv_triton', ['conv.cpp', 'conv.cu'], extra_cflags=['-O3'])