In this version, 2 main features were added:
- 1. GPU support
- The back-end GPU kernels are managed by the
gpu.js
library. - To add a Tensor to the gpu, pass 'gpu' as the
device
argument. - To add a whole layer (
nn.Linear
,nn.Block
,nn.MultiHeadSelfAttention
) to the gpu, , pass 'gpu' as thedevice
argument. - The syntax is just like the PyTorch counterpart.
- The back-end GPU kernels are managed by the
- 2. Saving and Loading models
- The functions are
torch.save
andtorch.load
. - The functionality is explained on the README.md file.
- The functions are