Limitations

Although RaNNCModel is designed to work in the manner of torch.nn.Module, it has the following limitations.

Control constructs

RaNNC uses a computation graph produced by PyTorch’s tracing function. As explained in PyTorch’s documentation, the tracing function does not record control constructs, including conditional branches and loops.

However, a function with @torch.jit.script can preserve control constructs even after tracing. You can call such a function from your model. test_function() in test/test_simple.py shows an example using a function with @torch.jit.script.

Arguments and return values

Arguments and outputs of RaNNCModel must satisfy the following conditions.

  • Arguments must be (mini-)batches tensors, whose first dimension corresponds to samples in a mini-batch.

  • Keyword arguments are not allowed.

  • Outputs must be (mini-)batches tensors or a loss value (scalar tensor).