Below is a script for training a custom LLM architecture, based on the GRU using Pytorch. I want you to implement code for saving the model to a file (checkpointing) and loading those weights if the file already exists  in the current working directory when initializing the model. The format should be a pytorch_model.bin format. You should add a variable: "save_every", similar to "infer_every", ...
Code:
```py
import torch
import math
from transformers import GPT2Tokenizer
from datasets import load_dataset
import numpy as np
import wandb  # Import W&B library

from model import minGRULM
from util import generate_text, generate_name



# ============================
# Configuration Parameters
# ============================
dataset_path  = 'flpelerin/tinystories-100k'
run_name      = generate_name()

num_epochs    = 1
batch_size    = 4
seq_length    = 256
learning_rate = 1e-4

input_len     = 50
num_predict   = 250

infer_every       = 200
reset_state_every = 16
validate_every    = 200



# ============================
# Initialize the Device
# ============================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Total context size is {batch_size * seq_length} tokens")



# ============================
# Initialize the Tokenizer
# ============================


tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token
vocab_size = tokenizer.vocab_size
print(f"Tokenzier has {vocab_size} unique tokens")





# ============================
# Load and Preprocess Dataset
# ============================
dataset = load_dataset(dataset_path)

def process_function(examples):
    return tokenizer(
        examples['text'],
        padding='max_length',  # Fixed padding
        truncation=True,
        max_length=(seq_length * batch_size)  # Fixed max length
    )

tokenized_datasets = dataset.map(process_function, batched=True)
print(f"Dataset has {tokenized_datasets['train'].num_rows} rows of {batch_size} times {seq_length} tokens")




# ============================
# Split Dataset into Train and Validation
# ============================
split_dataset = tokenized_datasets['train'].train_test_split(test_size=(1/validate_every))
train_dataset = split_dataset['train']
valid_dataset = split_dataset['test']

print(f"Training set size: {len(train_dataset)}")
print(f"Validation set size: {len(valid_dataset)}")



# ============================
# Initialize the Model
# ============================
model = minGRULM(
    vocab_size = vocab_size,
    d_model = 384,
    d_inner = 768,
    n_layers = 6
)

model.to(device)
parameters_count = sum(p.numel() for p in model.parameters())
print(f"Model has {parameters_count:,} parameters")



# ============================
# Initialize the Weights and Biases Run
# ============================
wandb.login(key="<TRUFFLEHOG_REDACTED_WEIGHTSANDBIASES>")
wandb.init(
    project="minGRU-Training",
    name=run_name,
    config={
        "dataset_path": dataset_path,
        
        "num_epochs": num_epochs,
        "batch_size": batch_size,
        "seq_length": seq_length,
        "learning_rate": learning_rate,

        "input_len": input_len,
        "num_predict": num_predict,

        "infer_every": infer_every,
        "reset_state_every": reset_state_every,
        "validate_every": validate_every,

        "dataset_rows": tokenized_datasets['train'].num_rows,
        "dataset_token_count": batch_size * seq_length,
        
        "train_set_size": len(train_dataset),
        "valid_set_size": len(valid_dataset),

        "model_parameters": parameters_count,

        "vocab_size": vocab_size,
        "d_model": 384,
        "d_inner": 768,
        "n_layers": 6,

        "device": str(device)
    }
)



# ============================
# Training Loop with Validation
# ============================
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
h_states = None
step = 0

for epoch in range(num_epochs):
    print(f"Starting Epoch {epoch + 1}/{num_epochs}")
    for i in range(0, len(train_dataset), batch_size):
        batch = train_dataset[i:i + batch_size]
        input_ids = torch.tensor(batch['input_ids']).to(device)

        # Reset hidden states if needed
        h_states = h_states if (step % reset_state_every!= 0) else None
        avg_states = sum([torch.mean(h_states[i]).item() for i in range(len(h_states))]) / len(h_states) if h_states is not None else None
        var_states = torch.var(torch.cat(h_states, dim=0)).item() if h_states else None

        optimizer.zero_grad()
        _, h_states, loss = model.forward(input_ids, h_states)
        loss.backward()
        optimizer.step()

        step += 1
        # Log step information
        wandb.log({
            "loss": loss.item(),
            "average_hidden_state": avg_states,
            "variance_hidden_state": var_states
        })
        print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden States: average = {avg_states}, variance = {var_states}")

        # Perform validation at specified intervals
        if step % validate_every == 0:
            validation_loss = 0.0
            valid_steps = 0

            with torch.no_grad():
                for vi in range(0, len(valid_dataset), batch_size):
                    val_batch = valid_dataset[vi:vi + batch_size]
                    val_input_ids = torch.tensor(val_batch['input_ids']).to(device)

                    # Forward pass
                    _, _, val_loss = model.forward(val_input_ids, None)
                    validation_loss += val_loss.item()
                    valid_steps += 1

            avg_validation_loss = validation_loss / valid_steps if valid_steps > 0 else float('inf')
            # Log validation loss
            wandb.log({"validation_loss": avg_validation_loss, "step": step})
            print(f"----- Validation after Step {step}: Average Loss = {avg_validation_loss:.4f} -----")

        # Perform inference at specified steps
        if step % infer_every == 0:
            with torch.no_grad():
                # Select a single input from the current batch for inference
                sample_ids = input_ids[0][:input_len]
                input_text = tokenizer.decode(sample_ids, skip_special_tokens=True)
                print(f"Input for Inference: {input_text}")

                prompt = sample_ids.unsqueeze(0)  # Shape: [1, input_len]
                generated_text = generate_text(model, tokenizer, prompt, num_predict)
                print(f"Generated Text:\n{generated_text}\n")
                # Optionally, log generated text (e.g., as HTML to preserve formatting)
                # wandb.log({"generated_text": wandb.Html(f"<pre>{generated_text}</pre>")}, step=step)
```