import torch
import torch.nn as nn
from torch.autograd import Variable

rnn = nn.RNNCell(10, 20)
gru = nn.GRUCell(10, 20)
input = Variable(torch.randn(3, 10))
h0 = Variable(torch.randn(1, 20)) # this should be (3, 20)
# both cells support broadcasting on CPU
hx_rnn = rnn(input, h0)
hx_gru = gru(input, h0)
assert hx_rnn.size() == hx_gru.size()

# send models and variables to GPU
rnn.cuda()
gru.cuda()
input = input.cuda()
h0 = h0.cuda()

# RNNCell is OK
try:
    hx_rnn_cuda = rnn(input, h0)
except RuntimeError as e:
    print('rnn error\n', e)

# GRUCell on GPU does not support broadcasting
try:
    hx_gru_cuda = gru(input, h0)
except RuntimeError as e:
    print('gru error\n', e)