Skip to content

Commit

Permalink
cleaning up docstrings
Browse files Browse the repository at this point in the history
  • Loading branch information
SamChou05 committed Dec 3, 2024
1 parent c58cd20 commit 8f74b54
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions afqinsight/nn/tests/test_pt_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def run_pytorch_model(
model,
device,
data_loaders,
n_epochs=100,
n_epochs=20,
permute=False,
):
"""
Expand All @@ -71,7 +71,7 @@ def run_pytorch_model(
Validation data loader
n_epochs : int
Number of epochs to train the model.
If no value provided, epochs is 100.
If no value provided, epochs is 20.
permute : boolean
Whether to permute the dimensions of the input batch for models
that require input with a specific shape.
Expand Down Expand Up @@ -183,7 +183,9 @@ def test_models(model, permute, device, data_loaders, data_shapes):
Testing data loader
Validation data loader
data_shapes : tuple
The shape of the input and target tensors.
The number of target classes,
Sequence length,
Number of input channels.
Returns
"""
Expand Down

0 comments on commit 8f74b54

Please sign in to comment.