Problem when using SNIP and Grasp for initialization

[dywu98]

It seems that when assigning --sparse_init GraSP there are problems:

Traceback (most recent call last):
File "/home/dancer/chase/Cifar/main_dst_chase.py", line 469, in
main()
File "/home/dancer/chase/Cifar/main_dst_chase.py", line 371, in main
mask.add_module(model, sparse_init=args.sparse_init, density=args.density)
File "/home/dancer/chase/Cifar/sparselearning/core_dst_chase.py", line 1508, in add_module
self.init(mode=sparse_init, density=density)
File "/home/dancer/chase/Cifar/sparselearning/core_dst_chase.py", line 1227, in init
layer_wise_sparsities = GraSP(self.module, self.density, self.train_loader, self.device)
File "/home/dancer/chase/Cifar/sparselearning/snip.py", line 217, in GraSP
inputs, targets = GraSP_fetch_data(train_dataloader, num_classes, samples_per_class)
File "/home/dancer/chase/Cifar/sparselearning/snip.py", line 160, in GraSP_fetch_data
if len(datas[category]) == samples_per_class:
IndexError: list index out of range

Moreover, when assigning --sparse_init SNIP, there is also a bug:

Traceback (most recent call last):
File "/home/dancer/chase/Cifar/main_dst_chase.py", line 469, in
main()
File "/home/dancer/chase/Cifar/main_dst_chase.py", line 395, in main
train(args, model, device, train_loader, epoch,mask)
File "/home/dancer/chase/Cifar/main_dst_chase.py", line 142, in train
if mask is not None: mask.step()
File "/home/dancer/chase/Cifar/sparselearning/core_dst_chase.py", line 1017, in step
self.truncate_weights_prune(self.death_rate)
File "/home/dancer/chase/Cifar/sparselearning/core_dst_chase.py", line 1786, in truncate_weights_prune
tokill=self.total_nonzero-self.baseline_nonzero
TypeError: unsupported operand type(s) for -: 'float' and 'NoneType'

So, could you please provide an instruction how to use other sparse init method aside from 'fixed_ERK'?