corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

HDF5 save and Conv2d stride typing change

Browse source
This commit is contained in:
Corentin Risselin 2020-07-04 13:05:07 +09:00
commit 58310ba077
1 changed files with 4 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

6
layers.py
View file

@ -1,3 +1,5 @@
from typing import Union, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
@ -44,8 +46,8 @@ class Layer(nn.Module):
class Conv2d(Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1,
activation=0, batch_norm=None, **kwargs):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
stride: Union[int, Tuple[int, int]] = 1, activation=0, batch_norm=None, **kwargs):
super().__init__(activation, batch_norm)
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)