pclucas14

from torch.autograd import Function, Variable
import numpy as np
import torch
import torch.nn as nn
import pdb
from torch.nn.modules.utils import _pair
from pyinn.utils import Dtype, Stream, load_kernel
import torch.nn.functional as F
from torch.autograd import gradcheck

pclucas14

class einsum_conv(nn.Module):
    def __init__(self, kernel_size, stride=1):
        super(einsum_conv, self).__init__()
        self.ks = kernel_size
        self.stride = stride

    def forward(self, x, kernel):
        if len(x.size()) == 3:
            x = x.unsqueeze(0)