[Python] What is the equivalent of torch.nn.Parameter(...) in julia's flux?

In pytorch I can create a custom module as follows (this code example is taken from here):

from torch import nn

class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.param = nn.Parameter(torch.randn(1, 1))

def forward(self, x):
x = x * self.param
return x

model = MyModel()
print(dict(model.named_parameters()))
# {'param': Parameter containing:
# tensor([[0.6077]], requires_grad=True)}

out = model(torch.randn(1, 1))
loss = out.mean()
loss.backward()

print(model.param.grad)
# tensor([[-1.3033]])


I want to be able to do the same using julia's flux. Specifically I am interested in knowing what the equivalent of nn.Parameter(torch.randn(1, 1)) is in flux.

P.S. I am tagging both python and julia as I believe this would increase the chance of this post reaching to people who has knowledge in both languages.

Continue reading...