How can I define a cell rather than a layer?

[GELIELEO]

[KhaleelKhan]

Hello @GELIELEO ,

Defining a Cell instead of a Layer is not possible with the interface we currently provide since this is a optimized implementation. Cell interface is currently under development.

[KhaleelKhan]

update: please check branch feature/egru_cell. Example usage is given below.

import torch
import evnn_pytorch as evnn

input_size = 128
hidden_size = 16
batch_size = 5

# setting use_custom_cuda=False makes the model use pytorch code instead of EvNN extension, 
# has no effect on the step function which is only implemented in pytorch
egru_cell =  evnn.EGRU(input_size, hidden_size, batch_first=False, use_custom_cuda=True)

egru_cell.cuda()

timesteps=25
# `x` is a CUDA tensor with shape [N,T,C]
x = torch.rand([timesteps, batch_size, input_size]).cuda()
y = []
y_t = None
h_t = None

for t in range(timesteps):
    x_t = x[t]
    y_t, h_t, _ = egru_cell.step(x_t, y_t, h_t)
    y.append(y_t)

y = torch.stack(y, dim=0)

# Sanity check, some numerical error from float conversion is expected 
assert torch.allclose(y, egru_cell(x)[0], atol=5e-4)