from stockroom import StockRoom
stock = StockRoom()

import torch

net = torch.nn.Sequential(
    torch.nn.Linear(10, 6),
    torch.nn.ReLU(),
    torch.nn.Linear(6, 3))

# storing state dict
stock.model['SeqNet'] = net.state_dict()

# or letting stockroom figure out how to fetch the state dict
stock.model.save('SeqNet', net)

# fetching state dict
state_dict = stock.model['SeqNet']
net.load_state_dict(state_dict)

# Letting stockroom does the state dict loading
stock.model.load('SeqNet', net)

import numpy as np

stock.data['arraysetname', 'samplename'] = np.array([1, 2])

# or

aset = stock.data['arraysetname']
aset['samplename'] = np.array([1, 2])

# on the CM open we'll make the checkout open and open the CM of the checkout
# for making the read and write optimized
with stock.optimize():
    for i, data in enumerage(get_data()):
        stock.data['arraysetname', i] = data

stock.tag['epochs'] = epochs
stock.tag['optimizer'] = 'SGD'
for e in range(epochs):
    for i, data in enumerate(get_batch_from_stockroom()):
        optimizer.zero_grad)_
        out = model(data)
        loss = loss_fn(out, target)
        loss.backward()
        optimizer.step()
        stock.tag[f'loss_{e}_{i}'] = loss.item()

import jovian
jovian.commit()