Merge branch 'master' of https://github.com/autonlab/DeepSurvivalMach…

…ines
autonlab · Feb 3, 2021 · d1930f6 · d1930f6
2 parents c03faf9 + 22b826a
commit d1930f6
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -15,3 +15,7 @@ install:
 script:
   - python -m pytest tests/
   - pylint --fail-under=8 dsm/
+
+#command to track code coverage
+after_success:
+  - bash <(curl -s https://codecov.io/bash)
diff --git a/dsm/__init__.py b/dsm/__init__.py
@@ -25,6 +25,8 @@
 r"""
 [![Build Status](https://travis-ci.org/autonlab/DeepSurvivalMachines.svg?branch=master)](https://travis-ci.org/autonlab/DeepSurvivalMachines)
 &nbsp;&nbsp;&nbsp;
+[![codecov](https://codecov.io/gh/autonlab/DeepSurvivalMachines/branch/main/graph/badge.svg?token=FU1HB5O92D)](https://codecov.io/gh/autonlab/DeepSurvivalMachines)
+&nbsp;&nbsp;&nbsp;
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 &nbsp;&nbsp;&nbsp;
 [![GitHub Repo stars](https://img.shields.io/github/stars/autonlab/DeepSurvivalMachines?style=social)](https://github.com/autonlab/DeepSurvivalMachines)

diff --git a/dsm/utilities.py b/dsm/utilities.py
@@ -57,7 +57,7 @@ def pretrain_dsm(model, t_train, e_train, t_valid, e_valid,
                                        risks=model.risks)
   premodel.double()
 
-  optimizer = torch.optim.Adam(premodel.parameters(), lr=lr)
+  optimizer = get_optimizer(model, lr)
 
   oldcost = float('inf')
   patience = 0
@@ -137,7 +137,7 @@ def train_dsm(model,
     model.scale[str(r+1)].data.fill_(float(premodel.scale[str(r+1)]))
 
   model.double()
-  optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+  optimizer = get_optimizer(model, lr)
 
   patience = 0
   oldcost = float('inf')